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