wangsenhao
85ce8c0e69
- Fix session resume: remove messages.length>1 guard so sessionStorage
ID is used as resumeSessionId on first message (prevents new session
creation when user sends "继续")
- Fix config_set tool schema: add missing `items: {}` to array type in
oneOf (Azure OpenAI strict validation rejects it otherwise)
- Metrics chart: use actual target set filename from trigger text
instead of generic "target_subset" key
- Add sub-agent skill call logging to agent_runtime.py (full
prompt/response, no truncation)
- Various metric enrichment and dedup fixes in server.py
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1373 lines
51 KiB
Python
1373 lines
51 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import re
|
|
import shlex
|
|
import ssl
|
|
import tarfile
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from html import unescape
|
|
from pathlib import Path, 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'
|
|
REMOTE_PIP_INDEX_URL = 'https://pypi.tuna.tsinghua.edu.cn/simple'
|
|
REMOTE_PIP_TRUSTED_HOST = 'pypi.tuna.tsinghua.edu.cn'
|
|
|
|
|
|
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()
|
|
account_part = sanitize_remote_path_part(binding.account_id)
|
|
self.account_runtime_root = (
|
|
f'{binding.workspace_root}/.runtime/accounts/{account_part}'
|
|
)
|
|
self.python_env_path = f'{self.account_runtime_root}/python/.venv'
|
|
self.platform_root = ''
|
|
self.skills_root = ''
|
|
self.bundle_digest = ''
|
|
self.synced_at: float | None = None
|
|
|
|
@classmethod
|
|
def from_persisted(cls, payload: dict[str, Any]) -> 'JupyterRuntimeSession':
|
|
"""从持久化的 Jupyter cookie 和绑定信息恢复 runtime。"""
|
|
|
|
if requests is None:
|
|
raise JupyterRuntimeError(
|
|
'Missing dependency: requests. Please install project dependencies.'
|
|
)
|
|
binding_payload = payload.get('binding')
|
|
if not isinstance(binding_payload, dict):
|
|
raise JupyterRuntimeError('Invalid persisted Jupyter binding.')
|
|
binding = JupyterWorkspaceBinding(
|
|
account_id=str(binding_payload['account_id']),
|
|
session_id=str(binding_payload['session_id']),
|
|
base_url=normalize_jupyter_base_url(str(binding_payload['base_url'])),
|
|
workspace_root=normalize_posix_path(str(binding_payload['workspace_root'])),
|
|
workspace_cwd=normalize_posix_path(str(binding_payload['workspace_cwd'])),
|
|
workspace_api_path=str(binding_payload['workspace_api_path']),
|
|
jupyter_tree_url=str(binding_payload['jupyter_tree_url']),
|
|
created_at=float(binding_payload.get('created_at') or time.time()),
|
|
)
|
|
http_session = requests.Session()
|
|
cookies = payload.get('cookies')
|
|
if isinstance(cookies, list):
|
|
for cookie in cookies:
|
|
if not isinstance(cookie, dict):
|
|
continue
|
|
name = cookie.get('name')
|
|
value = cookie.get('value')
|
|
if not isinstance(name, str) or not isinstance(value, str):
|
|
continue
|
|
http_session.cookies.set(
|
|
name,
|
|
value,
|
|
domain=(
|
|
cookie.get('domain')
|
|
if isinstance(cookie.get('domain'), str)
|
|
else None
|
|
),
|
|
path=(
|
|
cookie.get('path')
|
|
if isinstance(cookie.get('path'), str)
|
|
else '/'
|
|
),
|
|
)
|
|
xsrf_token = str(
|
|
payload.get('xsrf_token') or http_session.cookies.get('_xsrf') or ''
|
|
)
|
|
if not xsrf_token:
|
|
raise JupyterRuntimeError('Invalid persisted Jupyter _xsrf token.')
|
|
runtime = cls(
|
|
binding=binding,
|
|
http_session=http_session,
|
|
xsrf_token=xsrf_token,
|
|
)
|
|
if isinstance(payload.get('platform_root'), str):
|
|
runtime.platform_root = str(payload['platform_root'])
|
|
if isinstance(payload.get('skills_root'), str):
|
|
runtime.skills_root = str(payload['skills_root'])
|
|
if isinstance(payload.get('bundle_digest'), str):
|
|
runtime.bundle_digest = str(payload['bundle_digest'])
|
|
if isinstance(payload.get('synced_at'), (int, float)):
|
|
runtime.synced_at = float(payload['synced_at'])
|
|
return runtime
|
|
|
|
@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,
|
|
project_root: Path | None = None,
|
|
) -> '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,
|
|
project_root=project_root,
|
|
)
|
|
return runtime
|
|
|
|
def to_persisted_dict(self) -> dict[str, Any]:
|
|
"""导出可持久化的连接信息，用于页面刷新或服务重启后恢复。"""
|
|
|
|
return {
|
|
'binding': self.binding.to_dict(),
|
|
'xsrf_token': self._xsrf_token,
|
|
'cookies': [
|
|
{
|
|
'name': cookie.name,
|
|
'value': cookie.value,
|
|
'domain': cookie.domain,
|
|
'path': cookie.path,
|
|
'secure': bool(cookie.secure),
|
|
'expires': cookie.expires,
|
|
}
|
|
for cookie in self._session.cookies
|
|
],
|
|
'platform_root': self.platform_root,
|
|
'skills_root': self.skills_root,
|
|
'bundle_digest': self.bundle_digest,
|
|
'synced_at': self.synced_at,
|
|
'persisted_at': time.time(),
|
|
}
|
|
|
|
@property
|
|
def chat_workspace_root(self) -> str:
|
|
"""Per-user-per-chat root on shared NFS so the SFT training pod
|
|
(which doesn't mount the jupyter pod's private workspace) can read
|
|
and write the same artifacts. Layout:
|
|
/mnt/wangsenhao/autoresearch-zk-users/<account_id>/<session_id>/"""
|
|
account_part = sanitize_remote_path_part(self.binding.account_id)
|
|
session_part = sanitize_remote_path_part(self.binding.session_id)
|
|
return (
|
|
f'/mnt/wangsenhao/autoresearch-zk-users/'
|
|
f'{account_part}/{session_part}'
|
|
)
|
|
|
|
def bootstrap_workspace(
|
|
self,
|
|
*,
|
|
timeout_seconds: float = 20.0,
|
|
project_root: Path | None = None,
|
|
) -> None:
|
|
chat_root = self.chat_workspace_root
|
|
nfs_users_root = '/mnt/wangsenhao/autoresearch-zk-users'
|
|
probe = self.run_command(
|
|
(
|
|
f'mkdir -p {shlex.quote(nfs_users_root)} && '
|
|
f'touch {shlex.quote(nfs_users_root)}/.write_probe && '
|
|
f'rm -f {shlex.quote(nfs_users_root)}/.write_probe'
|
|
),
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=2000,
|
|
)
|
|
if probe.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
'远端 jupyter pod 未挂载 /mnt/wangsenhao 或不可写：'
|
|
+ (probe.stdout.strip() or probe.stderr.strip() or 'unknown error')
|
|
)
|
|
ws_output = f'{shlex.quote(self.binding.workspace_cwd)}/output'
|
|
chat_output_link = f'{shlex.quote(chat_root)}/output'
|
|
result = self.run_command(
|
|
(
|
|
'mkdir -p '
|
|
f'{ws_output} '
|
|
f'{shlex.quote(self.binding.workspace_cwd)}/input '
|
|
f'{shlex.quote(self.binding.workspace_cwd)}/scratchpad '
|
|
f'{shlex.quote(self.binding.workspace_root)}/.runtime/uploads '
|
|
f'{shlex.quote(self.binding.workspace_root)}/runtime_uploads '
|
|
f'{shlex.quote(self.account_runtime_root)}/python '
|
|
f'{shlex.quote(chat_root)}/results '
|
|
f'{shlex.quote(chat_root)}/sft_output '
|
|
f'{shlex.quote(chat_root)}/scripts && '
|
|
f'if [ -d {chat_output_link} ] && [ ! -L {chat_output_link} ]; then '
|
|
f'cp -a {chat_output_link}/* {ws_output}/ 2>/dev/null; '
|
|
f'rm -rf {chat_output_link}; '
|
|
f'fi && '
|
|
f'ln -sfnT {ws_output} {chat_output_link}'
|
|
),
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=4000,
|
|
)
|
|
if result.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
result.stdout.strip() or 'Unable to bootstrap Jupyter workspace.'
|
|
)
|
|
self.ensure_python_environment(timeout_seconds=max(timeout_seconds, 300.0))
|
|
if project_root is not None:
|
|
self.sync_runtime_bundle(project_root, timeout_seconds=max(timeout_seconds, 180.0))
|
|
self._sync_chat_workspace_scripts(timeout_seconds=timeout_seconds)
|
|
|
|
def _sync_chat_workspace_scripts(self, *, timeout_seconds: float = 20.0) -> None:
|
|
if not self.skills_root:
|
|
return
|
|
chat_root = self.chat_workspace_root
|
|
# SKILL/program §5 让 agent 直接 `bash scripts/<x>` 调下面这些脚本，
|
|
# 不同步过去 → agent 当 "脚本不存在" 处理后会肉手写 yaml，常把
|
|
# imageCommand 里的 `/scripts/prepare_and_train_sft.py` 写丢前缀，
|
|
# 导致训练 pod 报 "No such file or directory"。每次 bind 全量覆盖
|
|
# 一份，跟 prepare_and_train_sft.py 同样 "刷最新版"。
|
|
script_names = (
|
|
'prepare_and_train_sft.py',
|
|
'submit_sft.sh',
|
|
'sft_train_job.yaml.tpl',
|
|
'submit_cml_eval.sh',
|
|
'resolve_run_ids.sh',
|
|
)
|
|
src_dir = f'{self.skills_root}/model-iteration/scripts'
|
|
dst_dir = f'{chat_root}/scripts'
|
|
copy_cmds = []
|
|
for name in script_names:
|
|
src = f'{src_dir}/{name}'
|
|
dst = f'{dst_dir}/{name}'
|
|
copy_cmds.append(
|
|
f'if [ -f {shlex.quote(src)} ]; then '
|
|
f'cp {shlex.quote(src)} {shlex.quote(dst)} && '
|
|
f'chmod +x {shlex.quote(dst)} 2>/dev/null || true; '
|
|
f'fi'
|
|
)
|
|
result = self.run_command(
|
|
' && '.join(copy_cmds),
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=2000,
|
|
)
|
|
if result.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
result.stdout.strip() or 'Unable to sync chat workspace SFT script.'
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
payload = self.binding.to_dict()
|
|
payload['connected'] = True
|
|
payload['account_runtime_root'] = self.account_runtime_root
|
|
payload['python_env_path'] = self.python_env_path
|
|
payload['python_interpreter'] = self.python_interpreter
|
|
payload['platform_root'] = self.platform_root
|
|
payload['skills_root'] = self.skills_root
|
|
payload['bundle_digest'] = self.bundle_digest
|
|
payload['synced_at'] = self.synced_at
|
|
return payload
|
|
|
|
def render_context(self) -> str:
|
|
lines = [
|
|
'[远端 Jupyter 工作区]',
|
|
'当前 session 已切换到远端 Jupyter 运行时。',
|
|
f'- 默认工作目录：{self.binding.workspace_cwd}',
|
|
f'- 输出文件优先写入：{self.binding.workspace_cwd}/output',
|
|
f'- Python 环境：{self.python_env_path}',
|
|
'- Python 环境默认只初始化 venv 和 pip 清华源；缺包时使用 python_package 按需安装。',
|
|
'- 线上挂载目录（例如 /mnt/...）可以通过绝对路径读取和分析。',
|
|
]
|
|
if self.platform_root:
|
|
lines.extend(
|
|
[
|
|
f'- 平台运行包：{self.platform_root}',
|
|
f'- Skill 运行包：{self.skills_root}',
|
|
'- 读取 skills/...、src/...、标签定义/... 时会使用远端同步后的运行包。',
|
|
]
|
|
)
|
|
lines.extend(
|
|
[
|
|
'- 除非用户明确要求，不要修改平台代码目录；数据产物应留在当前 session 工作区。',
|
|
f'- Jupyter 工作区入口：{self.binding.jupyter_tree_url}',
|
|
]
|
|
)
|
|
return '\n'.join(lines)
|
|
|
|
@property
|
|
def python_interpreter(self) -> str:
|
|
return f'{self.python_env_path}/bin/python'
|
|
|
|
def ensure_python_environment(self, *, timeout_seconds: float = 300.0) -> None:
|
|
python_env = shlex.quote(self.python_env_path)
|
|
pip_conf = shlex.quote(f'{self.python_env_path}/pip.conf')
|
|
marker = shlex.quote(f'{self.python_env_path}/.zk-agent-python-env-v2')
|
|
command = f'''
|
|
set -e
|
|
if [ ! -x {python_env}/bin/python ]; then
|
|
python3 -m venv {python_env}
|
|
fi
|
|
cat > {pip_conf} <<'EOF'
|
|
[global]
|
|
index-url = {REMOTE_PIP_INDEX_URL}
|
|
trusted-host = {REMOTE_PIP_TRUSTED_HOST}
|
|
disable-pip-version-check = true
|
|
EOF
|
|
touch {marker}
|
|
{python_env}/bin/python -m pip --version
|
|
{python_env}/bin/python - <<'PY'
|
|
import sys
|
|
print(sys.executable)
|
|
PY
|
|
'''
|
|
result = self.run_command(
|
|
command,
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=12000,
|
|
)
|
|
if result.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
result.stdout.strip() or 'Unable to initialize remote Python env.'
|
|
)
|
|
|
|
def install_python_packages(
|
|
self,
|
|
packages: list[str],
|
|
*,
|
|
timeout_seconds: float,
|
|
max_output_chars: int,
|
|
cancel_event: Any | None = None,
|
|
) -> RemoteCommandResult:
|
|
package_args = ' '.join(shlex.quote(package) for package in packages)
|
|
return self.run_command(
|
|
f'{shlex.quote(self.python_interpreter)} -m pip install {package_args}',
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=max_output_chars,
|
|
cancel_event=cancel_event,
|
|
)
|
|
|
|
def sync_runtime_bundle(
|
|
self,
|
|
project_root: Path,
|
|
*,
|
|
timeout_seconds: float = 180.0,
|
|
) -> dict[str, Any]:
|
|
bundle_bytes, digest = build_runtime_bundle(project_root)
|
|
platform_root = f'{self.binding.workspace_root}/.runtime/platform/{digest[:16]}'
|
|
upload_path = f'{self.binding.workspace_root}/runtime_uploads/platform-{digest}.tar.gz'
|
|
digest_file = f'{platform_root}/.bundle_digest'
|
|
probe = self.run_command(
|
|
(
|
|
f'test -f {shlex.quote(digest_file)} '
|
|
f'&& cat {shlex.quote(digest_file)} || true'
|
|
),
|
|
timeout_seconds=20.0,
|
|
max_output_chars=2000,
|
|
)
|
|
if digest not in probe.stdout:
|
|
self.put_file_bytes(
|
|
api_path_for_absolute_path(upload_path),
|
|
bundle_bytes,
|
|
)
|
|
tmp_root = f'{platform_root}.tmp-{uuid.uuid4().hex[:8]}'
|
|
result = self.run_command(
|
|
(
|
|
f'rm -rf {shlex.quote(tmp_root)} && '
|
|
f'mkdir -p {shlex.quote(tmp_root)} && '
|
|
f'tar -xzf {shlex.quote(upload_path)} -C {shlex.quote(tmp_root)} && '
|
|
f'printf %s {shlex.quote(digest)} > {shlex.quote(tmp_root)}/.bundle_digest && '
|
|
f'rm -rf {shlex.quote(platform_root)} && '
|
|
f'mv {shlex.quote(tmp_root)} {shlex.quote(platform_root)}'
|
|
),
|
|
timeout_seconds=timeout_seconds,
|
|
max_output_chars=12000,
|
|
)
|
|
if result.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
result.stdout.strip() or 'Unable to extract remote runtime bundle.'
|
|
)
|
|
self.platform_root = platform_root
|
|
self.skills_root = f'{platform_root}/skills'
|
|
self.bundle_digest = digest
|
|
self.synced_at = time.time()
|
|
self.link_platform_resources()
|
|
return {
|
|
'bundle_digest': digest,
|
|
'platform_root': self.platform_root,
|
|
'skills_root': self.skills_root,
|
|
'uploaded_bytes': len(bundle_bytes),
|
|
}
|
|
|
|
def link_platform_resources(self) -> None:
|
|
if not self.platform_root:
|
|
return
|
|
workspace = shlex.quote(self.binding.workspace_cwd)
|
|
platform = shlex.quote(self.platform_root)
|
|
command = (
|
|
f'ln -sfn {platform}/skills {workspace}/skills && '
|
|
f'ln -sfn {platform}/src {workspace}/src && '
|
|
f'ln -sfn {platform}/标签定义 {workspace}/标签定义 && '
|
|
f'ln -sfn {platform}/pyproject.toml {workspace}/pyproject.toml'
|
|
)
|
|
result = self.run_command(
|
|
command,
|
|
timeout_seconds=20.0,
|
|
max_output_chars=4000,
|
|
)
|
|
if result.exit_code != 0:
|
|
raise JupyterRuntimeError(
|
|
result.stdout.strip() or 'Unable to link runtime resources into workspace.'
|
|
)
|
|
|
|
def put_file_bytes(self, api_path: str, data: bytes) -> None:
|
|
encoded = base64.b64encode(data).decode('ascii')
|
|
url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}'
|
|
response = self._session.put(
|
|
url,
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
json={
|
|
'type': 'file',
|
|
'format': 'base64',
|
|
'content': encoded,
|
|
},
|
|
timeout=60,
|
|
)
|
|
if response.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to upload file to Jupyter contents API: HTTP {response.status_code} {response.text[:500]}'
|
|
)
|
|
|
|
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)
|
|
env_prefix = self._remote_env_prefix()
|
|
inner_command = f'{env_prefix}{command}'
|
|
wrapped = (
|
|
f'mkdir -p {workspace} && '
|
|
f'cd {workspace} && '
|
|
f'bash -lc {shlex.quote(inner_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'{shlex.quote(self.python_interpreter)} {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 list_files(
|
|
self,
|
|
path: str,
|
|
*,
|
|
kind: str,
|
|
max_entries: int = 200,
|
|
max_output_chars: int = 20000,
|
|
) -> list[dict[str, Any]]:
|
|
target = self.resolve_workspace_path(path)
|
|
code = r'''
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
target = Path(PATH)
|
|
entries = []
|
|
if target.exists() and target.is_dir():
|
|
for item in sorted(target.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True)[:MAX_ENTRIES]:
|
|
if not item.is_file():
|
|
continue
|
|
stat = item.stat()
|
|
entries.append({
|
|
"name": item.name,
|
|
"path": str(item),
|
|
"kind": KIND,
|
|
"size": stat.st_size,
|
|
"modified_at": datetime.fromtimestamp(stat.st_mtime, timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
})
|
|
print(json.dumps(entries, ensure_ascii=False))
|
|
'''
|
|
script = (
|
|
code.replace('PATH', python_literal(target))
|
|
.replace('MAX_ENTRIES', str(max_entries))
|
|
.replace('KIND', python_literal(kind))
|
|
)
|
|
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_files failed')
|
|
try:
|
|
payload = json.loads(result.stdout.strip() or '[]')
|
|
except json.JSONDecodeError as exc:
|
|
raise JupyterRuntimeError('remote list_files returned invalid JSON') from exc
|
|
if not isinstance(payload, list):
|
|
raise JupyterRuntimeError('remote list_files returned invalid payload')
|
|
return [item for item in payload if isinstance(item, dict)]
|
|
|
|
def read_file_bytes(self, path: str) -> tuple[bytes, str]:
|
|
target = self.resolve_workspace_path(path)
|
|
api_path = api_path_for_absolute_path(target)
|
|
url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}'
|
|
response = self._session.get(
|
|
url,
|
|
params={'content': 1},
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
timeout=60,
|
|
)
|
|
if response.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to read remote file: HTTP {response.status_code} {response.text[:500]}'
|
|
)
|
|
payload = response.json()
|
|
if payload.get('type') != 'file':
|
|
raise JupyterRuntimeError('Remote path is not a file.')
|
|
content = payload.get('content')
|
|
file_format = payload.get('format')
|
|
if not isinstance(content, str):
|
|
raise JupyterRuntimeError('Remote file response did not contain content.')
|
|
if file_format == 'base64':
|
|
data = base64.b64decode(content)
|
|
else:
|
|
data = content.encode('utf-8')
|
|
name = str(payload.get('name') or PurePosixPath(target).name or 'remote-file')
|
|
return data, name
|
|
|
|
def file_info(self, path: str) -> dict[str, Any]:
|
|
target = self.resolve_workspace_path(path)
|
|
api_path = api_path_for_absolute_path(target)
|
|
url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}'
|
|
response = self._session.get(
|
|
url,
|
|
params={'content': 0},
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
timeout=30,
|
|
)
|
|
if response.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to inspect remote file: HTTP {response.status_code} {response.text[:500]}'
|
|
)
|
|
payload = response.json()
|
|
if payload.get('type') != 'file':
|
|
raise JupyterRuntimeError('Remote path is not a file.')
|
|
name = str(payload.get('name') or PurePosixPath(target).name or 'remote-file')
|
|
size = payload.get('size')
|
|
return {
|
|
'name': name,
|
|
'path': target,
|
|
'size': int(size) if isinstance(size, int) and size >= 0 else None,
|
|
}
|
|
|
|
def open_file_stream(self, path: str) -> tuple[Any, str]:
|
|
"""打开远端文件流，由调用方负责关闭 response。"""
|
|
|
|
target = self.resolve_workspace_path(path)
|
|
api_path = api_path_for_absolute_path(target)
|
|
url = f'{self.binding.base_url}/files/{quote(api_path, safe="/")}'
|
|
response = self._session.get(
|
|
url,
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
stream=True,
|
|
timeout=(10, 300),
|
|
)
|
|
if response.status_code >= 400:
|
|
response.close()
|
|
raise JupyterRuntimeError(
|
|
f'Unable to stream remote file: HTTP {response.status_code} {response.text[:500]}'
|
|
)
|
|
return response, str(PurePosixPath(target).name or 'remote-file')
|
|
|
|
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)
|
|
api_path = api_path_for_absolute_path(target)
|
|
if not api_path:
|
|
raise JupyterRuntimeError(f'Invalid remote path: {target}')
|
|
|
|
parent_api = '/'.join(api_path.split('/')[:-1])
|
|
if parent_api:
|
|
self._ensure_remote_dir(parent_api)
|
|
|
|
final_content = content
|
|
if append:
|
|
final_content = self._read_remote_text_for_append(api_path) + content
|
|
if newline_at_end and final_content and not final_content.endswith('\n'):
|
|
final_content += '\n'
|
|
|
|
encoded = base64.b64encode(final_content.encode('utf-8')).decode('ascii')
|
|
url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}'
|
|
response = self._session.put(
|
|
url,
|
|
json={
|
|
'type': 'file',
|
|
'format': 'base64',
|
|
'content': encoded,
|
|
},
|
|
headers={
|
|
'X-XSRFToken': self._xsrf_token,
|
|
'Content-Type': 'application/json',
|
|
},
|
|
timeout=60,
|
|
)
|
|
if response.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to write remote file: HTTP {response.status_code} {response.text[:500]}'
|
|
)
|
|
verb = 'appended' if append else 'wrote'
|
|
return f'{verb} {target} ({len(final_content)} chars)'
|
|
|
|
def _ensure_remote_dir(self, api_path: str) -> None:
|
|
parts = [segment for segment in api_path.split('/') if segment]
|
|
cumulative = ''
|
|
for segment in parts:
|
|
cumulative = f'{cumulative}/{segment}' if cumulative else segment
|
|
url = f'{self.binding.base_url}/api/contents/{quote(cumulative, safe="/")}'
|
|
resp = self._session.get(
|
|
url,
|
|
params={'content': 0},
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
timeout=30,
|
|
)
|
|
if resp.status_code == 200:
|
|
kind = resp.json().get('type')
|
|
if kind == 'directory':
|
|
continue
|
|
raise JupyterRuntimeError(
|
|
f'Cannot create directory at {cumulative}: path exists as {kind}'
|
|
)
|
|
if resp.status_code != 404:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to inspect remote path {cumulative}: HTTP {resp.status_code} {resp.text[:300]}'
|
|
)
|
|
create = self._session.put(
|
|
url,
|
|
json={'type': 'directory'},
|
|
headers={
|
|
'X-XSRFToken': self._xsrf_token,
|
|
'Content-Type': 'application/json',
|
|
},
|
|
timeout=30,
|
|
)
|
|
if create.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to create remote directory {cumulative}: HTTP {create.status_code} {create.text[:300]}'
|
|
)
|
|
|
|
def _read_remote_text_for_append(self, api_path: str) -> str:
|
|
url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}'
|
|
resp = self._session.get(
|
|
url,
|
|
params={'content': 1},
|
|
headers={'X-XSRFToken': self._xsrf_token},
|
|
timeout=60,
|
|
)
|
|
if resp.status_code == 404:
|
|
return ''
|
|
if resp.status_code >= 400:
|
|
raise JupyterRuntimeError(
|
|
f'Unable to read existing remote file for append: HTTP {resp.status_code} {resp.text[:500]}'
|
|
)
|
|
payload = resp.json()
|
|
if payload.get('type') != 'file':
|
|
raise JupyterRuntimeError(
|
|
f'Cannot append: remote path {api_path} is not a file (type={payload.get("type")})'
|
|
)
|
|
existing = payload.get('content', '')
|
|
if not isinstance(existing, str):
|
|
return ''
|
|
if payload.get('format') == 'base64':
|
|
return base64.b64decode(existing).decode('utf-8', errors='replace')
|
|
return existing
|
|
|
|
def resolve_workspace_path(self, raw_path: str) -> str:
|
|
value = raw_path.strip() or '.'
|
|
|
|
# 系统提示里给的"当前会话目录"通常是本地后端的绝对/相对路径，例如
|
|
# `.port_sessions/accounts/<acc>/sessions/<sid>/input` 或
|
|
# `/home/mi/.../.port_sessions/accounts/<acc>/sessions/<sid>/output/foo.csv`。
|
|
# 远端 Jupyter 没这套目录结构。识别这种 pattern 后，把 input/output/
|
|
# scratchpad bucket 映射到远端 workspace 的同名目录。
|
|
bucket_match = self._match_local_session_bucket(value)
|
|
if bucket_match is not None:
|
|
bucket, rest = bucket_match
|
|
tail = f'/{rest}' if rest else ''
|
|
return normalize_posix_path(
|
|
f'{self.binding.workspace_cwd}/{bucket}{tail}'
|
|
)
|
|
|
|
if value.startswith('/'):
|
|
return normalize_posix_path(value)
|
|
path = PurePosixPath(value)
|
|
if self.platform_root and path.parts:
|
|
platform_heads = {'skills', 'src', '标签定义'}
|
|
if path.parts[0] in platform_heads or value == 'pyproject.toml':
|
|
return normalize_posix_path(f'{self.platform_root}/{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}')
|
|
if path.parts and path.parts[0] in {'input', 'inputs'}:
|
|
tail = PurePosixPath(*path.parts[1:]) if len(path.parts) > 1 else PurePosixPath()
|
|
return normalize_posix_path(f'{self.binding.workspace_cwd}/input/{tail}')
|
|
return normalize_posix_path(f'{self.binding.workspace_cwd}/{value}')
|
|
|
|
_LOCAL_SESSION_BUCKET_RE = re.compile(
|
|
r'(?:^|/)\.port_sessions/accounts/[^/]+/sessions/[^/]+/'
|
|
r'(?P<bucket>input|inputs|output|outputs|scratchpad|scratch)'
|
|
r'(?:/(?P<rest>.*))?$'
|
|
)
|
|
|
|
@classmethod
|
|
def _match_local_session_bucket(cls, raw: str) -> tuple[str, str] | None:
|
|
m = cls._LOCAL_SESSION_BUCKET_RE.search(raw)
|
|
if not m:
|
|
return None
|
|
bucket = m.group('bucket')
|
|
normalized = {
|
|
'input': 'input',
|
|
'inputs': 'input',
|
|
'output': 'output',
|
|
'outputs': 'output',
|
|
'scratchpad': 'scratchpad',
|
|
'scratch': 'scratchpad',
|
|
}[bucket]
|
|
rest = m.group('rest') or ''
|
|
return normalized, rest
|
|
|
|
def _remote_env_prefix(self) -> str:
|
|
exports = {
|
|
'ZK_AGENT_WORKSPACE': self.binding.workspace_cwd,
|
|
'ZK_AGENT_OUTPUT': f'{self.binding.workspace_cwd}/output',
|
|
'ZK_AGENT_SCRATCHPAD': f'{self.binding.workspace_cwd}/scratchpad',
|
|
'ZK_AGENT_PYTHON_ENV': self.python_env_path,
|
|
'AUTORESEARCH_CHAT_ROOT': self.chat_workspace_root,
|
|
}
|
|
if self.platform_root:
|
|
exports['ZK_AGENT_PLATFORM_ROOT'] = self.platform_root
|
|
exports['ZK_AGENT_SKILLS_ROOT'] = self.skills_root
|
|
lines = [f'export {key}={shlex.quote(value)}' for key, value in exports.items()]
|
|
lines.append(f'export PATH={shlex.quote(self.python_env_path + "/bin")}:$PATH')
|
|
# CloudML 的 Jupyter Python 可能来自 Nix，pip 安装的科学计算轮子会依赖系统库。
|
|
# 只在库文件存在时预加载，避免 pandas/numpy/pyarrow 因 libz/libstdc++ 找不到而失败。
|
|
lines.extend(
|
|
[
|
|
'if [ -f /lib/x86_64-linux-gnu/libz.so.1 ] && '
|
|
'[ -f /lib/x86_64-linux-gnu/libstdc++.so.6 ]; then '
|
|
'export LD_PRELOAD=/lib/x86_64-linux-gnu/libz.so.1:'
|
|
'/lib/x86_64-linux-gnu/libstdc++.so.6${LD_PRELOAD:+:$LD_PRELOAD}; '
|
|
'fi'
|
|
]
|
|
)
|
|
if self.platform_root:
|
|
lines.append(f'export PYTHONPATH={shlex.quote(self.platform_root)}:${{PYTHONPATH:-}}')
|
|
return '; '.join(lines) + '; '
|
|
|
|
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,
|
|
project_root: Path | None = None,
|
|
) -> JupyterRuntimeSession:
|
|
runtime = JupyterRuntimeSession.connect(
|
|
account_id=account_id,
|
|
session_id=session_id,
|
|
base_url=base_url,
|
|
password=password,
|
|
workspace_root=workspace_root,
|
|
project_root=project_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 restore_session(
|
|
self,
|
|
*,
|
|
account_id: str,
|
|
session_id: str,
|
|
payload: dict[str, Any],
|
|
) -> JupyterRuntimeSession:
|
|
runtime = JupyterRuntimeSession.from_persisted(payload)
|
|
if (
|
|
runtime.binding.account_id != account_id
|
|
or runtime.binding.session_id != session_id
|
|
):
|
|
raise JupyterRuntimeError('Persisted Jupyter binding does not match session.')
|
|
with self._lock:
|
|
self._sessions[(account_id, session_id)] = runtime
|
|
return runtime
|
|
|
|
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 sync_account(self, account_id: str, project_root: Path) -> list[dict[str, Any]]:
|
|
with self._lock:
|
|
runtimes = [
|
|
runtime
|
|
for (runtime_account, _), runtime in self._sessions.items()
|
|
if runtime_account == account_id
|
|
]
|
|
results: list[dict[str, Any]] = []
|
|
for runtime in runtimes:
|
|
payload = runtime.sync_runtime_bundle(project_root)
|
|
results.append(
|
|
{
|
|
'session_id': runtime.binding.session_id,
|
|
**payload,
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
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 build_runtime_bundle(project_root: Path) -> tuple[bytes, str]:
|
|
root = project_root.resolve()
|
|
digest = calculate_runtime_bundle_digest(root)
|
|
buffer = io.BytesIO()
|
|
with tarfile.open(fileobj=buffer, mode='w:gz') as archive:
|
|
for name in ('src', 'skills', '标签定义', 'pyproject.toml'):
|
|
source = root / name
|
|
if not source.exists():
|
|
continue
|
|
archive.add(source, arcname=name, filter=_runtime_bundle_filter)
|
|
data = buffer.getvalue()
|
|
return data, digest
|
|
|
|
|
|
def calculate_runtime_bundle_digest(project_root: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
for name in ('src', 'skills', '标签定义', 'pyproject.toml'):
|
|
source = project_root / name
|
|
if not source.exists():
|
|
continue
|
|
if source.is_file():
|
|
_update_digest_for_file(digest, source, Path(name))
|
|
continue
|
|
for path in sorted(source.rglob('*')):
|
|
rel = Path(name) / path.relative_to(source)
|
|
if _should_skip_bundle_path(rel) or not path.is_file():
|
|
continue
|
|
_update_digest_for_file(digest, path, rel)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _update_digest_for_file(digest: 'hashlib._Hash', path: Path, rel: Path) -> None:
|
|
digest.update(rel.as_posix().encode('utf-8'))
|
|
digest.update(b'\0')
|
|
digest.update(path.read_bytes())
|
|
digest.update(b'\0')
|
|
|
|
|
|
def _should_skip_bundle_path(path: Path | PurePosixPath) -> bool:
|
|
parts = path.parts
|
|
skipped_dirs = {
|
|
'.git',
|
|
'.venv',
|
|
'.port_sessions',
|
|
'.next',
|
|
'node_modules',
|
|
'__pycache__',
|
|
}
|
|
if any(part in skipped_dirs for part in parts):
|
|
return True
|
|
return str(path).endswith(('.pyc', '.pyo', '.DS_Store'))
|
|
|
|
|
|
def _runtime_bundle_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None:
|
|
if _should_skip_bundle_path(PurePosixPath(info.name)):
|
|
return None
|
|
return info
|
|
|
|
|
|
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
|