diff --git a/backend/api/server.py b/backend/api/server.py index 3c49609..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 @@ -40,6 +41,11 @@ from src.agent_types import ( ) from src.bundled_skills import ALWAYS_ENABLED_HIDDEN_SKILL_NAMES, get_bundled_skills from src.data_agent_inputs import DataAgentInputError, load_input_sources +from src.jupyter_runtime import ( + DEFAULT_JUPYTER_WORKSPACE_ROOT, + JupyterRuntimeError, + JupyterRuntimeManager, +) from src.mcp_runtime import MCPRuntime, MCPServerProfile from src.openai_compat import OpenAICompatClient, OpenAICompatError from src.session_store import ( @@ -498,6 +504,7 @@ class AgentState: self._run_locks: dict[str, Lock] = {} self._account_configs: dict[str, AgentInstanceConfig] = {} self.run_manager = RunManager() + self.jupyter_runtime_manager = JupyterRuntimeManager() self._default_config = AgentInstanceConfig( cwd=cwd.resolve(), model=model, @@ -914,6 +921,14 @@ class ModelListRequest(BaseModel): account_id: str | None = None +class JupyterWorkspaceBindRequest(BaseModel): + session_id: str = Field(min_length=1) + base_url: str = Field(min_length=1) + password: str = Field(min_length=1) + workspace_root: str = DEFAULT_JUPYTER_WORKSPACE_ROOT + account_id: str | None = None + + class SkillPreferenceUpdate(BaseModel): skill: str | None = None enabled: bool @@ -940,6 +955,11 @@ class FeishuOnlineDocRequest(BaseModel): account_id: str | None = None +def _append_runtime_context(current: str | None, addition: str) -> str: + parts = [part.strip() for part in (current, addition) if part and part.strip()] + return '\n\n'.join(parts) + + # --------------------------------------------------------------------------- # App factory # --------------------------------------------------------------------------- @@ -971,6 +991,107 @@ def create_app(state: AgentState) -> FastAPI: raise HTTPException(status_code=400, detail=str(exc)) return state.snapshot(payload.account_id) + @app.get('/api/jupyter/session') + async def get_jupyter_session( + 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') + account_key = state._account_key(account_id) + return state.jupyter_runtime_manager.session_payload( + account_key, + 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, + ) -> dict[str, Any]: + safe_session_id = _safe_session_id(payload.session_id) + if not safe_session_id: + raise HTTPException(status_code=400, detail='session_id is required') + account_key = state._account_key(payload.account_id) + try: + runtime = await asyncio.to_thread( + state.jupyter_runtime_manager.bind_session, + account_id=account_key, + session_id=safe_session_id, + 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 + except Exception as exc: + raise HTTPException( + status_code=502, + detail=f'Unable to connect Jupyter runtime: {exc}', + ) from exc + return runtime.to_dict() + @app.get('/api/slash-commands') async def list_slash_commands() -> list[dict[str, Any]]: commands: list[dict[str, Any]] = [] @@ -1033,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') @@ -1068,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( @@ -1127,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, @@ -1312,6 +1445,16 @@ def create_app(state: AgentState) -> FastAPI: agent = state.agent_for(request.account_id, requested_session_id) config = state.config_for(request.account_id) session_directory = state.account_paths(request.account_id)['sessions'] + jupyter_runtime = state.jupyter_runtime_manager.get( + account_key, + requested_session_id, + ) + runtime_context = request.runtime_context + if jupyter_runtime is not None: + runtime_context = _append_runtime_context( + runtime_context, + jupyter_runtime.render_context(), + ) def emit_agent_event(event: dict[str, object]) -> None: stage = _runtime_event_stage(event) @@ -1371,6 +1514,7 @@ def create_app(state: AgentState) -> FastAPI: agent.tool_context, cancel_event=run_record.cancel_event, process_registry=run_record.process_registry, + jupyter_runtime=jupyter_runtime, ) started_at = time.perf_counter() try: @@ -1390,14 +1534,14 @@ def create_app(state: AgentState) -> FastAPI: result = agent.resume( prompt, stored, - runtime_context=request.runtime_context, + runtime_context=runtime_context, event_sink=emit_agent_event, ) else: result = agent.run( prompt, session_id=requested_session_id, - runtime_context=request.runtime_context, + runtime_context=runtime_context, event_sink=emit_agent_event, ) except Exception as exc: @@ -1426,6 +1570,7 @@ def create_app(state: AgentState) -> FastAPI: agent.tool_context, cancel_event=None, process_registry=None, + jupyter_runtime=None, ) elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000)) payload = _serialize_run_result(result) @@ -2094,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): @@ -2580,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', @@ -2596,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, } @@ -2616,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/app/api/claw/jupyter/route.ts b/frontend/app/app/api/claw/jupyter/route.ts new file mode 100644 index 0000000..aea48a4 --- /dev/null +++ b/frontend/app/app/api/claw/jupyter/route.ts @@ -0,0 +1,83 @@ +import { getCurrentAccount } from "@/lib/claw-auth"; + +const CLAW_API_URL = process.env.CLAW_API_URL ?? "http://127.0.0.1:8765"; + +type JupyterBindPayload = { + session_id?: string; + base_url?: string; + password?: string; + workspace_root?: string; +}; + +export async function GET(req: Request) { + const account = await getCurrentAccount(); + if (!account) + return Response.json({ error: "请先登录账号" }, { status: 401 }); + + const url = new URL(req.url); + const sessionId = url.searchParams.get("session_id"); + if (!sessionId) + return Response.json({ error: "缺少 session_id" }, { status: 400 }); + + try { + const target = new URL(`${CLAW_API_URL}/api/jupyter/session`); + target.searchParams.set("account_id", account.id); + target.searchParams.set("session_id", sessionId); + const response = await fetch(target, { cache: "no-store" }); + const body = await response.text(); + return new Response(body, { + status: response.status, + headers: { + "content-type": + response.headers.get("content-type") ?? "application/json", + }, + }); + } catch (err) { + return Response.json( + { + error: + err instanceof Error + ? `Unable to reach Claw backend: ${err.message}` + : "Unable to reach Claw backend", + }, + { status: 502 }, + ); + } +} + +export async function POST(req: Request) { + const account = await getCurrentAccount(); + if (!account) + return Response.json({ error: "请先登录账号" }, { status: 401 }); + + const payload = (await req.json()) as JupyterBindPayload; + try { + const response = await fetch(`${CLAW_API_URL}/api/jupyter/bind-session`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + ...payload, + account_id: account.id, + }), + cache: "no-store", + }); + const body = await response.text(); + return new Response(body, { + status: response.status, + headers: { + "content-type": + response.headers.get("content-type") ?? "application/json", + }, + }); + } catch (err) { + return Response.json( + { + error: + err instanceof Error + ? `Unable to reach Claw backend: ${err.message}` + : "Unable to reach Claw backend", + }, + { status: 502 }, + ); + } +} diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx index d7bb69a..9326d74 100644 --- a/frontend/app/components/assistant-ui/thread.tsx +++ b/frontend/app/components/assistant-ui/thread.tsx @@ -32,8 +32,22 @@ import { WrenchIcon, } from "lucide-react"; import { Popover as PopoverPrimitive } from "radix-ui"; -import type { ComponentProps, FC, KeyboardEvent, ReactNode } from "react"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import type { + ComponentProps, + Dispatch, + FC, + KeyboardEvent, + ReactNode, + SetStateAction, +} from "react"; +import { + useCallback, + useEffect, + useId, + useMemo, + useRef, + useState, +} from "react"; import TextareaAutosize from "react-textarea-autosize"; import { useActivityPanel } from "@/components/assistant-ui/activity-panel"; import { @@ -53,10 +67,13 @@ import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, + DialogDescription, + DialogFooter, DialogHeader, DialogTitle, DialogTrigger, } from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; import { Tooltip, TooltipContent, @@ -279,9 +296,13 @@ async function cancelLatestRun(sessionId: string, runId?: string | null) { const ChatTopActions: FC = () => { const { openFiles } = useActivityPanel(); + const [workspaceOpen, setWorkspaceOpen] = useState(false); const latestSessionId = useAuiState((s) => findLatestMessageMetadataValue(s.thread.messages, "sessionId"), ); + const normalizedSessionId = normalizeSessionId( + typeof latestSessionId === "string" ? latestSessionId : null, + ); return (
@@ -316,13 +337,269 @@ const ChatTopActions: FC = () => { 聊天中的文件 + +
); }; +type JupyterWorkspaceStatus = { + connected?: boolean; + session_id?: string; + workspace_cwd?: string; + jupyter_tree_url?: string; + detail?: string; + error?: string; +}; + +type WorkspaceSwitchProgress = { + percent: number; + label: string; +}; + +const WORKSPACE_SWITCH_STEPS = [ + "登录 Jupyter", + "初始化目录", + "创建 Python 环境", + "配置 pip 源", + "同步 Skill", + "链接工作区", +]; + +function WorkspaceSwitchDialog({ + open, + onOpenChange, + sessionId, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + sessionId: string | null; +}) { + const [jupyterUrl, setJupyterUrl] = useState(""); + const [password, setPassword] = useState(""); + const [workspaceRoot, setWorkspaceRoot] = useState( + "/root/zk_agent_workspaces", + ); + 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 ?? + (typeof window !== "undefined" + ? normalizeSessionId(readActiveSessionId()) + : null); + + useEffect(() => { + if (!open || !effectiveSessionId) return; + const sessionIdForRequest = effectiveSessionId; + let cancelled = false; + async function refreshStatus() { + const response = await fetch( + `/api/claw/jupyter?session_id=${encodeURIComponent(sessionIdForRequest)}`, + { cache: "no-store" }, + ); + const payload = (await response.json().catch(() => ({}))) as + | JupyterWorkspaceStatus + | Record; + if (!cancelled) setStatus(payload as JupyterWorkspaceStatus); + } + void refreshStatus().catch(() => undefined); + return () => { + cancelled = true; + }; + }, [open, effectiveSessionId]); + + async function submit() { + setMessage(""); + 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: sessionIdForRequest, + base_url: jupyterUrl.trim(), + password, + workspace_root: workspaceRoot.trim() || "/root/zk_agent_workspaces", + }), + }); + const payload = (await response.json().catch(() => ({}))) as + | JupyterWorkspaceStatus + | { detail?: string; error?: string }; + if (!response.ok) { + 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); + } + } + + return ( + + + + 切换工作区 + + 把当前 session 的工具执行切换到远端 Jupyter 开发机。 + + +
+ {status?.connected ? ( +
+
当前已连接
+
+ 工作目录：{status.workspace_cwd} +
+ {status.jupyter_tree_url ? ( + + 打开 Jupyter 工作区 + + ) : null} +
+ ) : null} + + + + {progress ? ( +
+
+ {progress.label} + + {Math.round(progress.percent)}% + +
+
+
+
+
+ ) : null} + {message ? ( +
+ {message} +
+ ) : null} +
+ + + + + +
+ ); +} + +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/pyproject.toml b/pyproject.toml index 00b7df7..896a0ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,8 @@ dependencies = [ "pydantic>=2.5", "openpyxl>=3.1", "pyarrow>=15", + "requests>=2.31", + "websocket-client>=1.7", ] [project.optional-dependencies] diff --git a/src/agent_tool_core.py b/src/agent_tool_core.py index 8ab9ecd..0ec8e0f 100644 --- a/src/agent_tool_core.py +++ b/src/agent_tool_core.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from .config_runtime import ConfigRuntime from .lsp_runtime import LSPRuntime from .mcp_runtime import MCPRuntime + from .jupyter_runtime import JupyterRuntimeSession from .plan_runtime import PlanRuntime from .remote_runtime import RemoteRuntime from .remote_trigger_runtime import RemoteTriggerRuntime @@ -49,6 +50,7 @@ class ToolExecutionContext: config_runtime: 'ConfigRuntime | None' = None lsp_runtime: 'LSPRuntime | None' = None mcp_runtime: 'MCPRuntime | None' = None + jupyter_runtime: 'JupyterRuntimeSession | None' = None remote_runtime: 'RemoteRuntime | None' = None remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None plan_runtime: 'PlanRuntime | None' = None @@ -138,6 +140,7 @@ def build_tool_context( config_runtime: 'ConfigRuntime | None' = None, lsp_runtime: 'LSPRuntime | None' = None, mcp_runtime: 'MCPRuntime | None' = None, + jupyter_runtime: 'JupyterRuntimeSession | None' = None, remote_runtime: 'RemoteRuntime | None' = None, remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None, plan_runtime: 'PlanRuntime | None' = None, @@ -169,6 +172,7 @@ def build_tool_context( config_runtime=config_runtime, lsp_runtime=lsp_runtime, mcp_runtime=mcp_runtime, + jupyter_runtime=jupyter_runtime, remote_runtime=remote_runtime, remote_trigger_runtime=remote_trigger_runtime, plan_runtime=plan_runtime, diff --git a/src/agent_tools.py b/src/agent_tools.py index fa752ad..b3d8de5 100644 --- a/src/agent_tools.py +++ b/src/agent_tools.py @@ -1683,6 +1683,8 @@ def _display_path(path: Path, context: ToolExecutionContext) -> str: def _execution_cwd(context: ToolExecutionContext) -> Path: + if context.jupyter_runtime is not None: + return Path(context.jupyter_runtime.binding.workspace_cwd) if context.scratchpad_directory is not None: context.scratchpad_directory.mkdir(parents=True, exist_ok=True) return context.scratchpad_directory @@ -1781,6 +1783,12 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str: if not isinstance(raw_path, str): raise ToolExecutionError('path must be a string') max_entries = _coerce_int(arguments, 'max_entries', 200) + if context.jupyter_runtime is not None: + return context.jupyter_runtime.list_dir( + raw_path, + max_entries=max_entries, + max_output_chars=context.max_output_chars, + ) target = _resolve_path(raw_path, context, allow_outside_root=True) if not target.exists(): raise ToolExecutionError(f'Path not found: {raw_path}') @@ -1798,6 +1806,23 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str: def _read_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: + if context.jupyter_runtime is not None: + start_line = arguments.get('start_line') + end_line = arguments.get('end_line') + if start_line is not None and ( + isinstance(start_line, bool) or not isinstance(start_line, int) or start_line < 1 + ): + raise ToolExecutionError('start_line must be an integer >= 1') + if end_line is not None and ( + isinstance(end_line, bool) or not isinstance(end_line, int) or end_line < 1 + ): + raise ToolExecutionError('end_line must be an integer >= 1') + return context.jupyter_runtime.read_text( + _require_string(arguments, 'path'), + start_line=start_line, + end_line=end_line, + max_output_chars=context.max_output_chars, + ) target = _resolve_path( _require_string(arguments, 'path'), context, @@ -1825,8 +1850,6 @@ def _read_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: _ensure_write_allowed(context) - target = _resolve_path(_require_string(arguments, 'path'), context) - _ensure_not_platform_code_write(target, context) content = _coerce_write_file_content(arguments) append = arguments.get('append', False) if not isinstance(append, bool): @@ -1834,6 +1857,25 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str newline_at_end = arguments.get('newline_at_end', False) if not isinstance(newline_at_end, bool): raise ToolExecutionError('newline_at_end must be a boolean') + if context.jupyter_runtime is not None: + message = context.jupyter_runtime.write_text( + _require_string(arguments, 'path'), + content, + append=append, + newline_at_end=newline_at_end, + max_output_chars=context.max_output_chars, + ) + return ( + message, + { + 'action': 'remote_append_file' if append else 'remote_write_file', + 'path': _require_string(arguments, 'path'), + 'remote_workspace': context.jupyter_runtime.binding.workspace_cwd, + 'content_length': len(content), + }, + ) + target = _resolve_path(_require_string(arguments, 'path'), context) + _ensure_not_platform_code_write(target, context) if newline_at_end and content and not content.endswith('\n'): content += '\n' previous_text: str | None = None @@ -2068,6 +2110,34 @@ def _notebook_edit(arguments: dict[str, Any], context: ToolExecutionContext) -> def _glob_search(arguments: dict[str, Any], context: ToolExecutionContext) -> str: pattern = _require_string(arguments, 'pattern') + if context.jupyter_runtime is not None: + code = r''' +from pathlib import Path +import glob +pattern = PATTERN +workspace = WORKSPACE +if pattern.startswith("/"): + matches = sorted(glob.glob(pattern, recursive=True)) +else: + matches = sorted(glob.glob(str(Path(workspace) / pattern), recursive=True)) +print("\n".join(matches) if matches else "(no matches)") +''' + script = ( + code.replace('PATTERN', repr(pattern)) + .replace('WORKSPACE', repr(context.jupyter_runtime.binding.workspace_cwd)) + ) + result = context.jupyter_runtime.run_python( + code=script, + script_path=None, + args=[], + stdin=None, + timeout_seconds=20.0, + max_output_chars=context.max_output_chars, + cancel_event=context.cancel_event, + ) + if result.exit_code != 0: + raise ToolExecutionError(result.stdout.strip() or 'remote glob_search failed') + return result.stdout.strip() or '(no matches)' matches = sorted(context.root.glob(pattern)) if not matches: return '(no matches)' @@ -2093,6 +2163,53 @@ def _grep_search(arguments: dict[str, Any], context: ToolExecutionContext) -> st if not isinstance(literal, bool): raise ToolExecutionError('literal must be a boolean') max_matches = _coerce_int(arguments, 'max_matches', 100) + if context.jupyter_runtime is not None: + code = r''' +from pathlib import Path +import re +pattern = PATTERN +literal = LITERAL +max_matches = MAX_MATCHES +root = Path(ROOT) +if not root.exists(): + raise SystemExit(f"Path not found: {root}") +regex = re.compile(re.escape(pattern) if literal else pattern) +hits = [] +files = root.rglob("*") if root.is_dir() else [root] +for file_path in files: + if not file_path.is_file(): + continue + try: + text = file_path.read_text(encoding="utf-8", errors="replace") + except OSError: + continue + for line_no, line in enumerate(text.splitlines(), start=1): + if regex.search(line): + hits.append(f"{file_path}:{line_no}: {line[:500]}") + if len(hits) >= max_matches: + print("\n".join(hits + [f"... truncated at {max_matches} matches ..."])) + raise SystemExit(0) +print("\n".join(hits) if hits else "(no matches)") +''' + root = context.jupyter_runtime.resolve_workspace_path(raw_path) + script = ( + code.replace('PATTERN', repr(pattern)) + .replace('LITERAL', repr(literal)) + .replace('MAX_MATCHES', str(max_matches)) + .replace('ROOT', repr(root)) + ) + result = context.jupyter_runtime.run_python( + code=script, + script_path=None, + args=[], + stdin=None, + timeout_seconds=30.0, + max_output_chars=context.max_output_chars, + cancel_event=context.cancel_event, + ) + if result.exit_code != 0: + raise ToolExecutionError(result.stdout.strip() or 'remote grep_search failed') + return result.stdout.strip() or '(no matches)' root = _resolve_path(raw_path, context, allow_outside_root=True) if not root.exists(): raise ToolExecutionError(f'Path not found: {raw_path}') @@ -2180,6 +2297,46 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) - context.command_timeout_seconds, ) max_output_chars = _coerce_int(arguments, 'max_output_chars', context.max_output_chars) + if context.jupyter_runtime is not None: + result = context.jupyter_runtime.run_python( + code=code, + script_path=( + context.jupyter_runtime.resolve_workspace_path(script_path) + if script_path + else None + ), + args=raw_args, + stdin=stdin, + timeout_seconds=timeout_seconds, + max_output_chars=max_output_chars, + cancel_event=context.cancel_event, + ) + payload = [ + f'exit_code={result.exit_code}', + f'interpreter={context.jupyter_runtime.python_interpreter}', + f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}', + '[stdout]', + result.stdout.rstrip(), + '[stderr]', + result.stderr.rstrip(), + ] + if result.timed_out: + payload.insert(0, f'timed_out=true\ntimeout_seconds={timeout_seconds:g}') + return ( + _truncate_output('\n'.join(payload).strip(), max_output_chars), + { + 'action': 'remote_python_exec', + 'mode': 'code' if code else 'script', + 'interpreter': context.jupyter_runtime.python_interpreter, + 'remote_workspace': context.jupyter_runtime.binding.workspace_cwd, + 'timed_out': result.timed_out, + 'timeout_seconds': timeout_seconds, + 'exit_code': result.exit_code, + 'stdout_preview': _snapshot_text(result.stdout), + 'stderr_preview': _snapshot_text(result.stderr), + 'output_preview': _snapshot_text('\n'.join(payload).strip()), + }, + ) interpreter = _resolve_python_interpreter(context) if code: command = [interpreter, '-c', code] @@ -2380,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': @@ -2484,6 +2692,34 @@ def _resolve_python_interpreter(context: ToolExecutionContext) -> str: def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str: command = _require_string(arguments, 'command') _ensure_shell_allowed(command, context) + if context.jupyter_runtime is not None: + result = context.jupyter_runtime.run_command( + command, + timeout_seconds=context.command_timeout_seconds, + max_output_chars=context.max_output_chars, + cancel_event=context.cancel_event, + ) + payload = [ + f'exit_code={result.exit_code}', + f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}', + '[stdout]', + result.stdout.rstrip(), + '[stderr]', + result.stderr.rstrip(), + ] + return ( + _truncate_output('\n'.join(payload).strip(), context.max_output_chars), + { + 'action': 'remote_bash', + 'command': command, + 'exit_code': result.exit_code, + 'remote_workspace': context.jupyter_runtime.binding.workspace_cwd, + 'timed_out': result.timed_out, + 'stdout_preview': _snapshot_text(result.stdout), + 'stderr_preview': _snapshot_text(result.stderr), + 'output_preview': _snapshot_text('\n'.join(payload).strip()), + }, + ) completed = subprocess.run( command, shell=True, @@ -3784,6 +4020,20 @@ def _stream_bash( arguments: dict[str, Any], context: ToolExecutionContext, ) -> Iterator[ToolStreamUpdate]: + if context.jupyter_runtime is not None: + result = _run_bash(arguments, context) + content, metadata = result if isinstance(result, tuple) else (result, {}) + tool_result = ToolExecutionResult( + name='bash', + ok=True, + content=content, + metadata=metadata, + ) + if content: + yield from _stream_static_text_result(tool_result) + else: + yield ToolStreamUpdate(kind='result', result=tool_result) + return try: command = _require_string(arguments, 'command') _ensure_shell_allowed(command, context) diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py new file mode 100644 index 0000000..e0f4d5a --- /dev/null +++ b/src/jupyter_runtime.py @@ -0,0 +1,1033 @@ +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 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 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)}/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, + ) + 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)) + + 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 read_text( + self, + path: str, + *, + start_line: int | None, + end_line: int | None, + max_output_chars: int, + ) -> str: + target = self.resolve_workspace_path(path) + code = r''' +from pathlib import Path +target = Path(PATH) +if not target.is_file(): + raise SystemExit(f"Path is not a file: {target}") +text = target.read_text(encoding="utf-8", errors="replace") +start_line = START_LINE +end_line = END_LINE +if start_line is None and end_line is None: + print(text, end="") +else: + lines = text.splitlines() + start_idx = max((start_line or 1) - 1, 0) + end_idx = end_line or len(lines) + print("\n".join(f"{start_idx + idx + 1}: {line}" for idx, line in enumerate(lines[start_idx:end_idx])), end="") +''' + script = ( + code.replace('PATH', python_literal(target)) + .replace('START_LINE', python_literal(start_line)) + .replace('END_LINE', python_literal(end_line)) + ) + result = self.run_python( + code=script, + script_path=None, + args=[], + stdin=None, + timeout_seconds=20.0, + max_output_chars=max_output_chars, + ) + if result.exit_code != 0: + raise JupyterRuntimeError(result.stdout.strip() or 'remote read_file failed') + return result.stdout[:max_output_chars] + + def write_text( + self, + path: str, + content: str, + *, + append: bool, + newline_at_end: bool, + max_output_chars: int, + ) -> str: + target = self.resolve_workspace_path(path) + if newline_at_end and content and not content.endswith('\n'): + content += '\n' + mode = 'a' if append else 'w' + code = r''' +from pathlib import Path +target = Path(PATH) +target.parent.mkdir(parents=True, exist_ok=True) +content = CONTENT +with target.open(MODE, encoding="utf-8") as fh: + fh.write(content) +print(f"{'appended' if MODE == 'a' else 'wrote'} {target} ({len(content)} chars)") +''' + script = ( + code.replace('PATH', python_literal(target)) + .replace('CONTENT', python_literal(content)) + .replace('MODE', python_literal(mode)) + ) + result = self.run_python( + code=script, + script_path=None, + args=[], + stdin=None, + timeout_seconds=20.0, + max_output_chars=max_output_chars, + ) + if result.exit_code != 0: + raise JupyterRuntimeError(result.stdout.strip() or 'remote write_file failed') + return result.stdout.strip() + + def resolve_workspace_path(self, raw_path: str) -> str: + value = raw_path.strip() or '.' + if value.startswith('/'): + return normalize_posix_path(value) + path = PurePosixPath(value) + if 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}') + 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 + 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 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