diff --git a/backend/api/server.py b/backend/api/server.py index 5683218..00f2fb7 100644 --- a/backend/api/server.py +++ b/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( - 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: diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx index 16d8717..f6cafe2 100644 --- a/frontend/app/components/assistant-ui/thread.tsx +++ b/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(() => + 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) { diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py index e0f4d5a..f1eb7c1 100644 --- a/src/jupyter_runtime.py +++ b/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: