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