diff --git a/backend/api/server.py b/backend/api/server.py index 644f8dd..cd76af3 100644 --- a/backend/api/server.py +++ b/backend/api/server.py @@ -23,10 +23,11 @@ from pathlib import Path from threading import Lock, RLock from typing import Any from urllib import error, request +from urllib.parse import quote, unquote, urlparse from uuid import uuid4 from fastapi import FastAPI, HTTPException -from fastapi.responses import FileResponse, JSONResponse, StreamingResponse +from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Field @@ -1004,6 +1005,66 @@ def create_app(state: AgentState) -> FastAPI: safe_session_id, ) + @app.get('/api/jupyter/files') + async def list_jupyter_files( + session_id: str, + account_id: str | None = None, + ) -> dict[str, Any]: + safe_session_id = _safe_session_id(session_id) + if not safe_session_id: + raise HTTPException(status_code=400, detail='session_id is required') + runtime = state.jupyter_runtime_manager.get( + state._account_key(account_id), + safe_session_id, + ) + if runtime is None: + return { + 'connected': False, + 'session_id': safe_session_id, + 'input': [], + 'output': [], + } + try: + return { + 'connected': True, + 'session_id': safe_session_id, + 'workspace_cwd': runtime.binding.workspace_cwd, + 'input': runtime.list_files('input', kind='input'), + 'output': runtime.list_files('output', kind='output'), + } + except JupyterRuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.get('/api/jupyter/file') + async def download_jupyter_file( + session_id: str, + path: str, + account_id: str | None = None, + ) -> Response: + safe_session_id = _safe_session_id(session_id) + if not safe_session_id: + raise HTTPException(status_code=400, detail='session_id is required') + runtime = state.jupyter_runtime_manager.get( + state._account_key(account_id), + safe_session_id, + ) + if runtime is None: + raise HTTPException(status_code=404, detail='Jupyter runtime is not connected') + try: + data, filename = runtime.read_file_bytes(path) + except JupyterRuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return Response( + data, + media_type=_content_type_for_filename(filename), + headers={ + 'content-disposition': ( + f'attachment; filename="{_ascii_download_filename(filename)}"; ' + f"filename*=UTF-8''{_url_quote_filename(filename)}" + ) + }, + ) + @app.post('/api/jupyter/bind-session') async def bind_jupyter_session( payload: JupyterWorkspaceBindRequest, @@ -1020,6 +1081,7 @@ def create_app(state: AgentState) -> FastAPI: base_url=payload.base_url, password=payload.password, workspace_root=payload.workspace_root, + project_root=state.config_for(payload.account_id).cwd, ) except JupyterRuntimeError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc @@ -1092,6 +1154,13 @@ def create_app(state: AgentState) -> FastAPI: except (ValueError, subprocess.TimeoutExpired) as exc: raise HTTPException(status_code=400, detail=str(exc)) result['skills'] = await list_skills(payload.account_id) + account_key = state._account_key(payload.account_id) + config = state.config_for(payload.account_id) + result['remote_runtime_sync'] = await asyncio.to_thread( + state.jupyter_runtime_manager.sync_account, + account_key, + config.cwd, + ) return result @app.get('/api/models') @@ -1127,7 +1196,11 @@ def create_app(state: AgentState) -> FastAPI: @app.post('/api/files/online-doc') async def create_feishu_online_doc(payload: FeishuOnlineDocRequest) -> dict[str, Any]: - file_path = _resolve_account_file(state, payload.account_id, payload.path) + file_path, online_doc_key = _resolve_feishu_source_file( + state, + payload.account_id, + payload.path, + ) suffix = file_path.suffix.lower() if suffix in FEISHU_UNSUPPORTED_DOCUMENT_SUFFIXES: raise HTTPException( @@ -1186,6 +1259,7 @@ def create_app(state: AgentState) -> FastAPI: state, payload.account_id, file_path=file_path, + map_key=online_doc_key, title=title, url=url, kind=kind, @@ -2165,6 +2239,28 @@ def _join_url(base_url: str, suffix: str) -> str: return f'{base_url.rstrip("/")}/{suffix.lstrip("/")}' +def _content_type_for_filename(filename: str) -> str: + suffix = Path(filename).suffix.lower() + if suffix == '.json': + return 'application/json; charset=utf-8' + if suffix in {'.jsonl', '.txt', '.md', '.csv'}: + return 'text/plain; charset=utf-8' + if suffix == '.xlsx': + return 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet' + if suffix == '.docx': + return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document' + return 'application/octet-stream' + + +def _ascii_download_filename(filename: str) -> str: + safe = re.sub(r'[^A-Za-z0-9._-]+', '_', filename).strip('._') + return safe[:120] or 'download' + + +def _url_quote_filename(filename: str) -> str: + return quote(filename, safe='') + + def _annotate_transcript_elapsed(payload: dict[str, Any], elapsed_ms: int) -> None: transcript = payload.get('transcript') if not isinstance(transcript, list): @@ -2651,6 +2747,7 @@ def _record_feishu_online_doc( account_id: str | None, *, file_path: Path, + map_key: str | None = None, title: str, url: str, kind: str = 'doc', @@ -2667,17 +2764,19 @@ def _record_feishu_online_doc( files = payload.get('files') if isinstance(payload, dict) else None if not isinstance(files, dict): files = {} - previous = files.get(str(file_path)) + key = map_key or str(file_path) + previous = files.get(key) created_at = ( previous.get('created_at') if isinstance(previous, dict) and isinstance(previous.get('created_at'), int) else now ) - files[str(file_path)] = { + files[key] = { 'url': clean_url, 'title': title, 'kind': kind, 'file_path': str(file_path), + 'source_path': key, 'created_at': created_at, 'updated_at': now, } @@ -2687,6 +2786,59 @@ def _record_feishu_online_doc( ) +def _resolve_feishu_source_file( + state: AgentState, + account_id: str | None, + raw_path: str, +) -> tuple[Path, str | None]: + if raw_path.startswith('jupyter://'): + return _materialize_jupyter_file_for_feishu(state, account_id, raw_path), raw_path + return _resolve_account_file(state, account_id, raw_path), None + + +def _materialize_jupyter_file_for_feishu( + state: AgentState, + account_id: str | None, + raw_uri: str, +) -> Path: + session_id, remote_path = _parse_jupyter_file_uri(raw_uri) + runtime = state.jupyter_runtime_manager.get( + state._account_key(account_id), + session_id, + ) + if runtime is None: + raise HTTPException(status_code=404, detail='Jupyter runtime is not connected') + try: + data, filename = runtime.read_file_bytes(remote_path) + except JupyterRuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + target_dir = ( + state.account_paths(account_id)['base'] + / 'integrations' + / 'feishu' + / 'remote-files' + / session_id + ) + target_dir.mkdir(parents=True, exist_ok=True) + target = target_dir / _safe_uploaded_filename(filename) + target.write_bytes(data) + return target + + +def _parse_jupyter_file_uri(raw_uri: str) -> tuple[str, str]: + parsed = urlparse(raw_uri) + session_id = _safe_session_id(parsed.netloc) + remote_path = unquote(parsed.path or '') + if not session_id or not remote_path.startswith('/'): + raise HTTPException(status_code=400, detail='Invalid Jupyter file URI') + return session_id, remote_path + + +def _safe_uploaded_filename(filename: str) -> str: + cleaned = re.sub(r'[\x00-\x1f/\\]+', '_', filename).strip(' ._') + return cleaned[:180] or 'remote-file' + + def _resolve_account_file(state: AgentState, account_id: str | None, raw_path: str) -> Path: path = Path(raw_path).expanduser().resolve() account_base = state.account_paths(account_id)['base'].resolve() diff --git a/frontend/app/app/api/claw/files/route.ts b/frontend/app/app/api/claw/files/route.ts index 6716966..94ffa9e 100644 --- a/frontend/app/app/api/claw/files/route.ts +++ b/frontend/app/app/api/claw/files/route.ts @@ -8,6 +8,7 @@ import { } from "@/lib/claw-auth"; export const runtime = "nodejs"; +const CLAW_API_URL = process.env.CLAW_API_URL ?? "http://127.0.0.1:8765"; export async function GET(req: Request) { const account = await getCurrentAccount(); @@ -19,15 +20,26 @@ export async function GET(req: Request) { const sessionId = normalizeSessionId(url.searchParams.get("session_id")); const onlineDocs = await readOnlineDocMap(account.id); if (!requestedPath && sessionId) { + const remoteFiles = await listRemoteSessionFiles( + account.id, + sessionId, + onlineDocs, + ); return Response.json({ session_id: sessionId, - input: await listSessionFiles(account.id, sessionId, "input", onlineDocs), - output: await listSessionFiles( - account.id, - sessionId, - "output", - onlineDocs, - ), + input: [ + ...(await listSessionFiles(account.id, sessionId, "input", onlineDocs)), + ...remoteFiles.input, + ], + output: [ + ...(await listSessionFiles( + account.id, + sessionId, + "output", + onlineDocs, + )), + ...remoteFiles.output, + ], }); } if (!requestedPath) { @@ -35,23 +47,38 @@ export async function GET(req: Request) { if (!latestSessionId) { return Response.json({ session_id: null, input: [], output: [] }); } + const remoteFiles = await listRemoteSessionFiles( + account.id, + latestSessionId, + onlineDocs, + ); return Response.json({ session_id: latestSessionId, - input: await listSessionFiles( - account.id, - latestSessionId, - "input", - onlineDocs, - ), - output: await listSessionFiles( - account.id, - latestSessionId, - "output", - onlineDocs, - ), + input: [ + ...(await listSessionFiles( + account.id, + latestSessionId, + "input", + onlineDocs, + )), + ...remoteFiles.input, + ], + output: [ + ...(await listSessionFiles( + account.id, + latestSessionId, + "output", + onlineDocs, + )), + ...remoteFiles.output, + ], }); } + if (requestedPath.startsWith("jupyter://")) { + return downloadRemoteFile(account.id, requestedPath); + } + const accountRoot = path.resolve(accountBaseRoot(account.id)); const filePath = path.resolve(requestedPath); if (!filePath.startsWith(`${accountRoot}${path.sep}`)) { @@ -75,6 +102,123 @@ export async function GET(req: Request) { } } +async function listRemoteSessionFiles( + accountId: string, + sessionId: string, + onlineDocs: Record, +) { + try { + const target = new URL(`${CLAW_API_URL}/api/jupyter/files`); + target.searchParams.set("account_id", accountId); + target.searchParams.set("session_id", sessionId); + const response = await fetch(target, { cache: "no-store" }); + if (!response.ok) return { input: [], output: [] }; + const payload = (await response.json()) as { + input?: RemoteSessionFile[]; + output?: RemoteSessionFile[]; + }; + return { + input: mapRemoteFiles(payload.input, sessionId, "input", onlineDocs), + output: mapRemoteFiles(payload.output, sessionId, "output", onlineDocs), + }; + } catch { + return { input: [], output: [] }; + } +} + +type RemoteSessionFile = { + name?: unknown; + path?: unknown; + kind?: unknown; + size?: unknown; + modified_at?: unknown; +}; + +function mapRemoteFiles( + files: RemoteSessionFile[] | undefined, + sessionId: string, + kind: "input" | "output", + onlineDocs: Record, +) { + if (!Array.isArray(files)) return []; + return files + .map((file) => { + const remotePath = + typeof file.path === "string" && file.path.startsWith("/") + ? file.path + : null; + if (!remotePath) return null; + const name = + typeof file.name === "string" && file.name.trim() + ? file.name.trim() + : path.basename(remotePath); + const uri = toJupyterFileUri(sessionId, remotePath); + const onlineDoc = onlineDocs[uri]; + return { + name, + path: uri, + kind, + source: "jupyter", + size: typeof file.size === "number" ? file.size : 0, + modified_at: + typeof file.modified_at === "string" + ? file.modified_at + : new Date().toISOString(), + download_url: `/api/claw/files?path=${encodeURIComponent(uri)}`, + online_doc_url: cleanOnlineDocUrl(onlineDoc?.url) ?? null, + online_doc_title: onlineDoc?.title ?? null, + online_doc_kind: onlineDoc?.kind ?? null, + online_doc_updated_at: onlineDoc?.updated_at ?? null, + }; + }) + .filter((file): file is NonNullable => Boolean(file)); +} + +async function downloadRemoteFile(accountId: string, uri: string) { + const parsed = parseJupyterFileUri(uri); + if (!parsed) { + return Response.json({ error: "无效的远端文件路径" }, { status: 400 }); + } + const target = new URL(`${CLAW_API_URL}/api/jupyter/file`); + target.searchParams.set("account_id", accountId); + target.searchParams.set("session_id", parsed.sessionId); + target.searchParams.set("path", parsed.remotePath); + const response = await fetch(target, { cache: "no-store" }); + if (!response.ok) { + const payload = await response.text(); + return new Response(payload, { + status: response.status, + headers: { + "content-type": + response.headers.get("content-type") ?? "application/json", + }, + }); + } + return new Response(response.body, { + status: response.status, + headers: { + "content-type": + response.headers.get("content-type") ?? "application/octet-stream", + "content-disposition": + response.headers.get("content-disposition") ?? + `attachment; filename="${encodeURIComponent(path.basename(parsed.remotePath))}"`, + }, + }); +} + +function toJupyterFileUri(sessionId: string, remotePath: string) { + return `jupyter://${sessionId}${remotePath}`; +} + +function parseJupyterFileUri(uri: string) { + const match = uri.match(/^jupyter:\/\/([^/]+)(\/.*)$/u); + if (!match) return null; + return { + sessionId: match[1], + remotePath: match[2], + }; +} + async function listSessionFiles( accountId: string, sessionId: string, diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx index 0d77551..9326d74 100644 --- a/frontend/app/components/assistant-ui/thread.tsx +++ b/frontend/app/components/assistant-ui/thread.tsx @@ -32,7 +32,14 @@ import { WrenchIcon, } from "lucide-react"; import { Popover as PopoverPrimitive } from "radix-ui"; -import type { ComponentProps, FC, KeyboardEvent, ReactNode } from "react"; +import type { + ComponentProps, + Dispatch, + FC, + KeyboardEvent, + ReactNode, + SetStateAction, +} from "react"; import { useCallback, useEffect, @@ -359,6 +366,20 @@ type JupyterWorkspaceStatus = { error?: string; }; +type WorkspaceSwitchProgress = { + percent: number; + label: string; +}; + +const WORKSPACE_SWITCH_STEPS = [ + "登录 Jupyter", + "初始化目录", + "创建 Python 环境", + "配置 pip 源", + "同步 Skill", + "链接工作区", +]; + function WorkspaceSwitchDialog({ open, onOpenChange, @@ -376,6 +397,9 @@ function WorkspaceSwitchDialog({ const [status, setStatus] = useState(null); const [submitting, setSubmitting] = useState(false); const [message, setMessage] = useState(""); + const [progress, setProgress] = useState( + null, + ); const fieldId = useId(); const effectiveSessionId = sessionId ?? @@ -405,21 +429,19 @@ function WorkspaceSwitchDialog({ async function submit() { setMessage(""); - if (!effectiveSessionId) { - setMessage("请先发送一条消息创建 session，再切换工作区。"); - return; - } + const sessionIdForRequest = ensureWorkspaceSessionId(effectiveSessionId); if (!jupyterUrl.trim() || !password) { setMessage("请填写 Jupyter 地址和密码。"); return; } setSubmitting(true); + const progressTimer = startWorkspaceProgress(setProgress); try { const response = await fetch("/api/claw/jupyter", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ - session_id: effectiveSessionId, + session_id: sessionIdForRequest, base_url: jupyterUrl.trim(), password, workspace_root: workspaceRoot.trim() || "/root/zk_agent_workspaces", @@ -432,11 +454,15 @@ function WorkspaceSwitchDialog({ setMessage(payload.detail ?? payload.error ?? "切换工作区失败"); return; } + writeActiveSessionId(sessionIdForRequest); setStatus(payload as JupyterWorkspaceStatus); + setProgress({ percent: 100, label: "切换完成" }); setMessage("工作区切换成功，后续工具会在远端 Jupyter 工作区执行。"); setPassword(""); } finally { + window.clearInterval(progressTimer); setSubmitting(false); + window.setTimeout(() => setProgress(null), 1200); } } @@ -475,6 +501,7 @@ function WorkspaceSwitchDialog({ value={jupyterUrl} onChange={(event) => setJupyterUrl(event.target.value)} placeholder="https://...-jupyter.../lab" + disabled={submitting} /> + {progress ? ( +
+
+ {progress.label} + + {Math.round(progress.percent)}% + +
+
+
+
+
+ ) : null} {message ? (
{message} @@ -518,6 +563,43 @@ function WorkspaceSwitchDialog({ ); } +function ensureWorkspaceSessionId(current: string | null) { + if (current) return current; + const generated = + typeof crypto !== "undefined" && "randomUUID" in crypto + ? `__LOCALID_${crypto.randomUUID().replace(/-/g, "").slice(0, 8)}` + : `__LOCALID_${Math.random().toString(36).slice(2, 10)}`; + writeActiveSessionId(generated); + scheduleSessionListRefresh(); + return generated; +} + +function startWorkspaceProgress( + setProgress: Dispatch>, +) { + setProgress({ percent: 6, label: WORKSPACE_SWITCH_STEPS[0] }); + return window.setInterval(() => { + setProgress((current) => { + const previous = current ?? { + percent: 6, + label: WORKSPACE_SWITCH_STEPS[0], + }; + const nextPercent = Math.min( + 92, + previous.percent + (previous.percent < 55 ? 9 : 4), + ); + const stepIndex = Math.min( + WORKSPACE_SWITCH_STEPS.length - 1, + Math.floor((nextPercent / 100) * WORKSPACE_SWITCH_STEPS.length), + ); + return { + percent: nextPercent, + label: WORKSPACE_SWITCH_STEPS[stepIndex], + }; + }); + }, 1200); +} + const ThreadMessage: FC = () => { const role = useAuiState((s) => s.message.role); const isEditing = useAuiState((s) => s.message.composer.isEditing); diff --git a/src/agent_tools.py b/src/agent_tools.py index 753be9d..b3d8de5 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -2313,7 +2313,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) - ) payload = [ f'exit_code={result.exit_code}', - 'interpreter=remote:python3', + f'interpreter={context.jupyter_runtime.python_interpreter}', f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}', '[stdout]', result.stdout.rstrip(), @@ -2327,7 +2327,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) - { 'action': 'remote_python_exec', 'mode': 'code' if code else 'script', - 'interpreter': 'remote:python3', + 'interpreter': context.jupyter_runtime.python_interpreter, 'remote_workspace': context.jupyter_runtime.binding.workspace_cwd, 'timed_out': result.timed_out, 'timeout_seconds': timeout_seconds, @@ -2537,13 +2537,64 @@ def _communicate_after_stop(process: subprocess.Popen[str]) -> tuple[str, str]: def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str: _ensure_process_execution_allowed(context, 'Python package management') action = _require_string(arguments, 'action') - interpreter = _resolve_python_interpreter(context) timeout_seconds = _coerce_float( arguments, 'timeout_seconds', min(context.command_timeout_seconds * 4, 120.0), ) max_output_chars = _coerce_int(arguments, 'max_output_chars', context.max_output_chars) + if context.jupyter_runtime is not None: + interpreter = context.jupyter_runtime.python_interpreter + if action == 'show': + result = context.jupyter_runtime.run_command( + f'{shlex.quote(interpreter)} -m pip --version', + timeout_seconds=timeout_seconds, + max_output_chars=max_output_chars, + cancel_event=context.cancel_event, + ) + elif action == 'install': + packages = arguments.get('packages') + if not isinstance(packages, list) or not packages: + raise ToolExecutionError('packages must be a non-empty array for action=install') + package_args = [str(package).strip() for package in packages if str(package).strip()] + if not package_args: + raise ToolExecutionError('packages must contain at least one non-empty package name') + result = context.jupyter_runtime.install_python_packages( + package_args, + timeout_seconds=timeout_seconds, + max_output_chars=max_output_chars, + cancel_event=context.cancel_event, + ) + else: + raise ToolExecutionError('action must be "show" or "install"') + payload = [ + f'exit_code={result.exit_code}', + f'interpreter={interpreter}', + f'python_env={context.jupyter_runtime.python_env_path}', + '[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_package', + 'package_action': action, + 'interpreter': interpreter, + 'python_env_dir': context.jupyter_runtime.python_env_path, + '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 action == 'show': command = [interpreter, '-m', 'pip', '--version'] elif action == 'install': diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py index af9ceaf..e0f4d5a 100644 --- a/src/jupyter_runtime.py +++ b/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,