Add Jupyter workspace runtime

backend/api/server.py

@@ -40,6 +40,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 +503,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 +920,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 +954,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 +990,46 @@ 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.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,
+            )
+        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]] = []
@@ -1312,6 +1371,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 +1440,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 +1460,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 +1496,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)