wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge Jupyter workspace runtime

Browse Source
This commit is contained in:
wuyang6
2026-05-13 11:00:23 +08:00
parent 126f35ae46 86f49134a4
commit b0579a1247
8 changed files with 2046 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+229 -6
View File
@@ -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()
frontend/app/app/api/claw/files/route.ts
+163 -19
View File
@@ -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<string, OnlineDocRecord>,
) {
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<string, OnlineDocRecord>,
) {
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<typeof file> => 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,
frontend/app/app/api/claw/jupyter/route.ts
+83
View File
@@ -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 },
);
}
}
frontend/app/components/assistant-ui/thread.tsx
+279 -2
View File
@@ -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 (
<div className="pointer-events-none sticky top-3 z-20 mr-4 flex justify-end">
@@ -316,13 +337,269 @@ const ChatTopActions: FC = () => {
<FileTextIcon className="size-4 text-muted-foreground" />
聊天中的文件
</button>
<button
type="button"
className="flex w-full items-center gap-2 rounded-lg px-3 py-2 text-left text-sm transition-colors hover:bg-muted"
onClick={() => setWorkspaceOpen(true)}
>
<WrenchIcon className="size-4 text-muted-foreground" />
切换工作区
</button>
</PopoverPrimitive.Content>
</PopoverPrimitive.Portal>
</PopoverPrimitive.Root>
<WorkspaceSwitchDialog
open={workspaceOpen}
onOpenChange={setWorkspaceOpen}
sessionId={normalizedSessionId}
/>
</div>
);
};
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<JupyterWorkspaceStatus | null>(null);
const [submitting, setSubmitting] = useState(false);
const [message, setMessage] = useState("");
const [progress, setProgress] = useState<WorkspaceSwitchProgress | null>(
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<string, never>;
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 (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="pointer-events-auto sm:max-w-xl">
<DialogHeader>
<DialogTitle>切换工作区</DialogTitle>
<DialogDescription>
把当前 session 的工具执行切换到远端 Jupyter 开发机。
</DialogDescription>
</DialogHeader>
<div className="grid gap-4 text-sm">
{status?.connected ? (
<div className="rounded-md border bg-muted/40 p-3 text-muted-foreground">
<div className="font-medium text-foreground">当前已连接</div>
<div className="mt-1 break-all">
工作目录：{status.workspace_cwd}
</div>
{status.jupyter_tree_url ? (
<a
className="mt-1 block break-all text-primary underline-offset-4 hover:underline"
href={status.jupyter_tree_url}
target="_blank"
rel="noreferrer"
>
打开 Jupyter 工作区
</a>
) : null}
</div>
) : null}
<label className="grid gap-1.5" htmlFor={`${fieldId}-url`}>
<span className="text-muted-foreground">Jupyter 地址</span>
<Input
id={`${fieldId}-url`}
value={jupyterUrl}
onChange={(event) => setJupyterUrl(event.target.value)}
placeholder="https://...-jupyter.../lab"
disabled={submitting}
/>
</label>
<label className="grid gap-1.5" htmlFor={`${fieldId}-password`}>
<span className="text-muted-foreground">密码</span>
<Input
id={`${fieldId}-password`}
type="password"
value={password}
onChange={(event) => setPassword(event.target.value)}
placeholder="Jupyter 登录密码"
disabled={submitting}
/>
</label>
<label className="grid gap-1.5" htmlFor={`${fieldId}-root`}>
<span className="text-muted-foreground">远端工作区根目录</span>
<Input
id={`${fieldId}-root`}
value={workspaceRoot}
onChange={(event) => setWorkspaceRoot(event.target.value)}
disabled={submitting}
/>
</label>
{progress ? (
<div className="rounded-md border bg-muted/30 px-3 py-2">
<div className="mb-2 flex items-center justify-between gap-3 text-xs">
<span className="text-muted-foreground">{progress.label}</span>
<span className="font-medium tabular-nums">
{Math.round(progress.percent)}%
</span>
</div>
<div className="h-1.5 overflow-hidden rounded-full bg-muted">
<div
className="h-full rounded-full bg-primary transition-all duration-500"
style={{ width: `${progress.percent}%` }}
/>
</div>
</div>
) : null}
{message ? (
<div className="rounded-md bg-muted px-3 py-2 text-muted-foreground">
{message}
</div>
) : null}
</div>
<DialogFooter>
<Button
type="button"
variant="outline"
onClick={() => onOpenChange(false)}
>
关闭
</Button>
<Button type="button" disabled={submitting} onClick={submit}>
{submitting ? "连接中..." : "切换"}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
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<SetStateAction<WorkspaceSwitchProgress | null>>,
) {
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);
pyproject.toml
+2
View File
@@ -38,6 +38,8 @@ dependencies = [
"pydantic>=2.5",
"openpyxl>=3.1",
"pyarrow>=15",
"requests>=2.31",
"websocket-client>=1.7",
]
[project.optional-dependencies]
src/agent_tool_core.py
+4
View File
@@ -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,
src/agent_tools.py
+253 -3
View File
@@ -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)
src/jupyter_runtime.py
+1033
View File
File diff suppressed because it is too large Load Diff