Persist Jupyter workspace bindings

backend/api/server.py

@@ -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(
+        runtime = _jupyter_runtime_for_session(
-            account_key,
+            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(
+        runtime = _jupyter_runtime_for_session(state, account_id, safe_session_id)
-            state._account_key(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(
+        runtime = _jupyter_runtime_for_session(state, account_id, safe_session_id)
-            state._account_key(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(
+        jupyter_runtime = _jupyter_runtime_for_session(
-            account_key,
+            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(
+    runtime = _jupyter_runtime_for_session(state, account_id, session_id)
-        state._account_key(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

@@ -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

@@ -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: