Improve Jupyter workspace bootstrap

src/jupyter_runtime.py

@@ -1,16 +1,19 @@
 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 PurePosixPath
+from pathlib import Path, PurePosixPath
 from typing import Any
 from urllib.parse import quote, urlparse, urlunparse
 
@@ -26,6 +29,8 @@ except ImportError:  # pragma: no cover
 
 
 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):
@@ -79,6 +84,15 @@ class JupyterRuntimeSession:
         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 connect(
@@ -90,6 +104,7 @@ class JupyterRuntimeSession:
         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(
@@ -123,15 +138,27 @@ class JupyterRuntimeSession:
             http_session=http_session,
             xsrf_token=xsrf_token,
         )
-        runtime.bootstrap_workspace(timeout_seconds=timeout_seconds)
+        runtime.bootstrap_workspace(
+            timeout_seconds=timeout_seconds,
+            project_root=project_root,
+        )
         return runtime
 
-    def bootstrap_workspace(self, *, timeout_seconds: float = 20.0) -> None:
+    def bootstrap_workspace(
+        self,
+        *,
+        timeout_seconds: float = 20.0,
+        project_root: Path | None = None,
+    ) -> None:
         result = self.run_command(
             (
                 'mkdir -p '
                 f'{shlex.quote(self.binding.workspace_cwd)}/output '
-                f'{shlex.quote(self.binding.workspace_cwd)}/scratchpad'
+                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'
             ),
             timeout_seconds=timeout_seconds,
             max_output_chars=4000,
@@ -140,24 +167,190 @@ class JupyterRuntimeSession:
             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))
 
     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:
-        return '\n'.join(
+        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(
             [
-                '[远端 Jupyter 工作区]',
-                '当前 session 已切换到远端 Jupyter 运行时。',
-                f'- 默认工作目录：{self.binding.workspace_cwd}',
-                f'- 输出文件优先写入：{self.binding.workspace_cwd}/output',
-                '- 线上挂载目录（例如 /mnt/...）可以通过绝对路径读取和分析。',
                 '- 除非用户明确要求，不要修改平台代码目录；数据产物应留在当前 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,
@@ -175,10 +368,12 @@ class JupyterRuntimeSession:
             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(command)}; '
+            f'bash -lc {shlex.quote(inner_command)}; '
             f'printf "\\n{marker}:%s\\n" "$?"'
         )
         with self._lock:
@@ -235,7 +430,7 @@ class JupyterRuntimeSession:
             )
         assert script_path is not None
         quoted_args = ' '.join(shlex.quote(arg) for arg in args)
-        python_command = f'python3 {shlex.quote(script_path)}'
+        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:
@@ -283,6 +478,86 @@ print("\n".join(entries) if entries else "(empty directory)")
             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 read_text(
         self,
         path: str,
@@ -369,6 +644,10 @@ print(f"{'appended' if MODE == 'a' else 'wrote'} {target} ({len(content)} chars)
         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}')
@@ -377,6 +656,33 @@ print(f"{'appended' if MODE == 'a' else 'wrote'} {target} ({len(content)} chars)
             return normalize_posix_path(f'{self.binding.workspace_cwd}/scratchpad/{tail}')
         return normalize_posix_path(f'{self.binding.workspace_cwd}/{value}')
 
+    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,
+        }
+        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
@@ -500,6 +806,7 @@ class JupyterRuntimeManager:
         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,
@@ -507,6 +814,7 @@ class JupyterRuntimeManager:
             base_url=base_url,
             password=password,
             workspace_root=workspace_root,
+            project_root=project_root,
         )
         with self._lock:
             self._sessions[(account_id, session_id)] = runtime
@@ -526,6 +834,24 @@ class JupyterRuntimeManager:
             }
         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()
@@ -601,6 +927,65 @@ 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,