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

Persist Jupyter workspace bindings

Browse Source
This commit is contained in:
wuyang6
2026-05-13 17:59:50 +08:00
parent b7dba5f5b8
commit ef703ff49b
3 changed files with 208 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+67 -16
View File
@@ -46,6 +46,7 @@ from src.jupyter_runtime import (
DEFAULT_JUPYTER_WORKSPACE_ROOT,
JupyterRuntimeError,
JupyterRuntimeManager,
JupyterRuntimeSession,
)
from src.mcp_runtime import MCPRuntime, MCPServerProfile
from src.openai_compat import OpenAICompatClient, OpenAICompatError
@@ -81,6 +82,7 @@ FEISHU_SPREADSHEET_MAX_COLS = 200
FEISHU_SPREADSHEET_MAX_CELL_CHARS = 5000
FEISHU_SPREADSHEET_WRITE_BATCH_ROWS = 500
FEISHU_ONLINE_DOCS_FILENAME = 'online-docs.json'
JUPYTER_WORKSPACE_FILENAME = 'jupyter_workspace.json'
@dataclass
@@ -1041,10 +1043,18 @@ def create_app(state: AgentState) -> FastAPI:
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,
runtime = _jupyter_runtime_for_session(
state,
account_id,
safe_session_id,
)
if runtime is None:
return {
'connected': False,
'account_id': account_key,
'session_id': safe_session_id,
}
return runtime.to_dict()
@app.get('/api/jupyter/files')
async def list_jupyter_files(
@@ -1054,10 +1064,7 @@ def create_app(state: AgentState) -> FastAPI:
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,
)
runtime = _jupyter_runtime_for_session(state, account_id, safe_session_id)
if runtime is None:
return {
'connected': False,
@@ -1085,10 +1092,7 @@ def create_app(state: AgentState) -> FastAPI:
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,
)
runtime = _jupyter_runtime_for_session(state, account_id, safe_session_id)
if runtime is None:
raise HTTPException(status_code=404, detail='Jupyter runtime is not connected')
try:
@@ -1131,6 +1135,10 @@ def create_app(state: AgentState) -> FastAPI:
status_code=502,
detail=f'Unable to connect Jupyter runtime: {exc}',
) from exc
_save_jupyter_runtime_binding(
state.account_paths(payload.account_id)['sessions'],
runtime,
)
return runtime.to_dict()
@app.get('/api/slash-commands')
@@ -1615,8 +1623,9 @@ 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,
jupyter_runtime = _jupyter_runtime_for_session(
state,
request.account_id,
requested_session_id,
)
runtime_context = request.runtime_context
@@ -2654,6 +2663,51 @@ def _session_json_path(directory: Path, session_id: str) -> Path:
return directory / f'{session_id}.json'
def _jupyter_binding_path(directory: Path, session_id: str) -> Path:
return directory / session_id / JUPYTER_WORKSPACE_FILENAME
def _save_jupyter_runtime_binding(
directory: Path,
runtime: JupyterRuntimeSession,
) -> None:
path = _jupyter_binding_path(directory, runtime.binding.session_id)
try:
path.parent.mkdir(parents=True, exist_ok=True)
_write_session_metadata(path, runtime.to_persisted_dict())
except OSError:
return
def _jupyter_runtime_for_session(
state: AgentState,
account_id: str | None,
session_id: str,
) -> JupyterRuntimeSession | None:
account_key = state._account_key(account_id)
runtime = state.jupyter_runtime_manager.get(account_key, session_id)
if runtime is not None:
return runtime
path = _jupyter_binding_path(
state.account_paths(account_id)['sessions'],
session_id,
)
try:
payload = json.loads(path.read_text(encoding='utf-8'))
except (OSError, json.JSONDecodeError):
return None
if not isinstance(payload, dict):
return None
try:
return state.jupyter_runtime_manager.restore_session(
account_id=account_key,
session_id=session_id,
payload=payload,
)
except JupyterRuntimeError:
return None
def _feishu_paths(state: AgentState, account_id: str | None) -> dict[str, Path]:
base = state.account_paths(account_id)['base'] / 'integrations' / 'feishu'
paths = {
@@ -2992,10 +3046,7 @@ def _materialize_jupyter_file_for_feishu(
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,
)
runtime = _jupyter_runtime_for_session(state, account_id, session_id)
if runtime is None:
raise HTTPException(status_code=404, detail='Jupyter runtime is not connected')
try:
frontend/app/components/assistant-ui/thread.tsx
+35 -2
View File
@@ -80,6 +80,7 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import {
ACTIVE_SESSION_CHANGED_EVENT,
clearActiveSessionId,
clearPendingWorkspaceSessionId,
PENDING_WORKSPACE_SESSION_CHANGED_EVENT,
@@ -310,15 +311,47 @@ function usePendingWorkspaceSessionId() {
return sessionId;
}
function useCurrentThreadSessionId({ includePending = false } = {}) {
function useActiveSessionId() {
const [sessionId, setSessionId] = useState<string | null>(() =>
typeof window !== "undefined"
? normalizeSessionId(readActiveSessionId())
: null,
);
useEffect(() => {
const refresh = () => {
setSessionId(normalizeSessionId(readActiveSessionId()));
};
window.addEventListener(ACTIVE_SESSION_CHANGED_EVENT, refresh);
window.addEventListener("storage", refresh);
window.addEventListener("focus", refresh);
return () => {
window.removeEventListener(ACTIVE_SESSION_CHANGED_EVENT, refresh);
window.removeEventListener("storage", refresh);
window.removeEventListener("focus", refresh);
};
}, []);
return sessionId;
}
function useCurrentThreadSessionId({
includePending = false,
includeActive = true,
} = {}) {
const currentRun = useReplayedRunState();
const pendingSessionId = usePendingWorkspaceSessionId();
const activeSessionId = useActiveSessionId();
useEffect(() => {
if (currentRun.sessionId && currentRun.sessionId === pendingSessionId) {
clearPendingWorkspaceSessionId();
}
}, [currentRun.sessionId, pendingSessionId]);
return currentRun.sessionId ?? (includePending ? pendingSessionId : null);
return (
currentRun.sessionId ??
(includePending ? pendingSessionId : null) ??
(includeActive ? activeSessionId : null)
);
}
async function cancelLatestRun(sessionId: string, runId?: string | null) {
src/jupyter_runtime.py
+106
View File
@@ -94,6 +94,71 @@ class JupyterRuntimeSession:
self.bundle_digest = ''
self.synced_at: float | None = None
@classmethod
def from_persisted(cls, payload: dict[str, Any]) -> 'JupyterRuntimeSession':
"""从持久化的 Jupyter cookie 和绑定信息恢复 runtime。"""
if requests is None:
raise JupyterRuntimeError(
'Missing dependency: requests. Please install project dependencies.'
)
binding_payload = payload.get('binding')
if not isinstance(binding_payload, dict):
raise JupyterRuntimeError('Invalid persisted Jupyter binding.')
binding = JupyterWorkspaceBinding(
account_id=str(binding_payload['account_id']),
session_id=str(binding_payload['session_id']),
base_url=normalize_jupyter_base_url(str(binding_payload['base_url'])),
workspace_root=normalize_posix_path(str(binding_payload['workspace_root'])),
workspace_cwd=normalize_posix_path(str(binding_payload['workspace_cwd'])),
workspace_api_path=str(binding_payload['workspace_api_path']),
jupyter_tree_url=str(binding_payload['jupyter_tree_url']),
created_at=float(binding_payload.get('created_at') or time.time()),
)
http_session = requests.Session()
cookies = payload.get('cookies')
if isinstance(cookies, list):
for cookie in cookies:
if not isinstance(cookie, dict):
continue
name = cookie.get('name')
value = cookie.get('value')
if not isinstance(name, str) or not isinstance(value, str):
continue
http_session.cookies.set(
name,
value,
domain=(
cookie.get('domain')
if isinstance(cookie.get('domain'), str)
else None
),
path=(
cookie.get('path')
if isinstance(cookie.get('path'), str)
else '/'
),
)
xsrf_token = str(
payload.get('xsrf_token') or http_session.cookies.get('_xsrf') or ''
)
if not xsrf_token:
raise JupyterRuntimeError('Invalid persisted Jupyter _xsrf token.')
runtime = cls(
binding=binding,
http_session=http_session,
xsrf_token=xsrf_token,
)
if isinstance(payload.get('platform_root'), str):
runtime.platform_root = str(payload['platform_root'])
if isinstance(payload.get('skills_root'), str):
runtime.skills_root = str(payload['skills_root'])
if isinstance(payload.get('bundle_digest'), str):
runtime.bundle_digest = str(payload['bundle_digest'])
if isinstance(payload.get('synced_at'), (int, float)):
runtime.synced_at = float(payload['synced_at'])
return runtime
@classmethod
def connect(
cls,
@@ -144,6 +209,30 @@ class JupyterRuntimeSession:
)
return runtime
def to_persisted_dict(self) -> dict[str, Any]:
"""导出可持久化的连接信息，用于页面刷新或服务重启后恢复。"""
return {
'binding': self.binding.to_dict(),
'xsrf_token': self._xsrf_token,
'cookies': [
{
'name': cookie.name,
'value': cookie.value,
'domain': cookie.domain,
'path': cookie.path,
'secure': bool(cookie.secure),
'expires': cookie.expires,
}
for cookie in self._session.cookies
],
'platform_root': self.platform_root,
'skills_root': self.skills_root,
'bundle_digest': self.bundle_digest,
'synced_at': self.synced_at,
'persisted_at': time.time(),
}
def bootstrap_workspace(
self,
*,
@@ -824,6 +913,23 @@ class JupyterRuntimeManager:
with self._lock:
return self._sessions.get((account_id, session_id))
def restore_session(
self,
*,
account_id: str,
session_id: str,
payload: dict[str, Any],
) -> JupyterRuntimeSession:
runtime = JupyterRuntimeSession.from_persisted(payload)
if (
runtime.binding.account_id != account_id
or runtime.binding.session_id != session_id
):
raise JupyterRuntimeError('Persisted Jupyter binding does not match session.')
with self._lock:
self._sessions[(account_id, session_id)] = runtime
return runtime
def session_payload(self, account_id: str, session_id: str) -> dict[str, Any]:
runtime = self.get(account_id, session_id)
if runtime is None: