diff --git a/scribe/notebook/_notebook_server_utils.py b/scribe/notebook/_notebook_server_utils.py index 4a0fd12..d34d80a 100644 --- a/scribe/notebook/_notebook_server_utils.py +++ b/scribe/notebook/_notebook_server_utils.py @@ -7,6 +7,7 @@ import requests from fastmcp.utilities.types import Image +from jupyter_client.kernelspec import NoSuchKernel from ._image_processing_utils import resize_image_if_needed @@ -54,6 +55,39 @@ def clean_notebook_for_save(nb): return nb +def get_notebook_metadata_for_kernel(kernel_spec_manager, kernel_name: str) -> Dict[str, Any]: + """Build notebook metadata by querying the installed kernel spec. + + Args: + kernel_spec_manager: Jupyter's KernelSpecManager instance (from ServerApp). + kernel_name: Registered kernel name (e.g. "python3", "ir"). + + Returns: + Dict with "kernelspec" and stub "language_info" suitable for nb.metadata.update(). + + Raises: + ValueError: If the kernel is not installed. + """ + try: + spec = kernel_spec_manager.get_kernel_spec(kernel_name) + except NoSuchKernel: + available = ", ".join(sorted(kernel_spec_manager.find_kernel_specs().keys())) + raise ValueError( + f"Kernel '{kernel_name}' is not installed. Available kernels: {available}" + ) + + return { + "kernelspec": { + "display_name": spec.display_name, + "language": spec.language, + "name": kernel_name, + }, + # Stub language_info from the spec; the full version is populated + # after kernel startup via _get_kernel_language_info(). + "language_info": {"name": spec.language}, + } + + def check_server_health(port: int) -> Optional[Dict[str, Any]]: """Check if scribe server is running on given port.""" try: diff --git a/scribe/notebook/notebook_mcp_server.py b/scribe/notebook/notebook_mcp_server.py index 1848b6a..d873d58 100644 --- a/scribe/notebook/notebook_mcp_server.py +++ b/scribe/notebook/notebook_mcp_server.py @@ -148,6 +148,7 @@ async def _start_session_internal( notebook_path: Optional[str] = None, fork_prev_notebook: bool = True, tool_name: str = "start_session", + kernel_name: str = "python3", ) -> Dict[str, Any]: """ Internal helper function for starting sessions from scratch versus resuming versus forking existing notebook. @@ -157,13 +158,14 @@ async def _start_session_internal( notebook_path: Path to existing notebook (if any) fork_prev_notebook: If True, create new notebook; if False, use existing in-place tool_name: Name of the calling tool for logging/debugging + kernel_name: Jupyter kernel spec name (e.g. "python3", "ir") """ try: # Ensure server is running server_url = ensure_server_running() # Build request body - request_body = {} + request_body = {"kernel_name": kernel_name} if experiment_name: request_body["experiment_name"] = experiment_name if notebook_path: @@ -187,9 +189,8 @@ async def _start_session_internal( "status": "started", "notebook_path": data["notebook_path"], "vscode_url": f"{data.get('server_url', server_url)}/?token={data.get('token', token)}", - "kernel_name": data.get( - "kernel_name", data.get("kernel_display_name", "Scribe Kernel") - ), + "kernel_name": data.get("kernel_name", "python3"), + "kernel_display_name": data.get("kernel_display_name", "Scribe Kernel"), } # Track session for cleanup @@ -263,12 +264,18 @@ async def _start_session_internal( @mcp.tool -async def start_new_session(experiment_name: Optional[str] = None) -> Dict[str, Any]: +async def start_new_session( + experiment_name: Optional[str] = None, + kernel_name: str = "python3", +) -> Dict[str, Any]: """ Start a completely new Jupyter kernel session with an empty notebook. + Supports any installed Jupyter kernel (e.g. "python3", "ir" for R via IRkernel). + Args: experiment_name: Custom name for the notebook (e.g., "ImageGeneration") + kernel_name: Jupyter kernel spec name. Defaults to "python3". Returns: Dictionary with: @@ -284,6 +291,7 @@ async def start_new_session(experiment_name: Optional[str] = None) -> Dict[str, notebook_path=None, fork_prev_notebook=True, tool_name="start_new_session", + kernel_name=kernel_name, ) @@ -295,7 +303,9 @@ async def start_new_session(experiment_name: Optional[str] = None) -> Dict[str, "title": "Start Session - Resume Notebook" # A human-readable title for the tool. }, ) -async def start_session_resume_notebook(notebook_path: str) -> Dict[str, Any]: +async def start_session_resume_notebook( + notebook_path: str, kernel_name: str = "python3", +) -> Dict[str, Any]: """ Start a new session by resuming an existing notebook in-place, modifying the original notebook file. @@ -304,6 +314,7 @@ async def start_session_resume_notebook(notebook_path: str) -> Dict[str, Any]: Args: notebook_path: Path to the existing notebook to resume from + kernel_name: Jupyter kernel spec name. Should match the notebook's language. Returns: Dictionary with: @@ -322,12 +333,14 @@ async def start_session_resume_notebook(notebook_path: str) -> Dict[str, Any]: notebook_path=notebook_path, fork_prev_notebook=False, tool_name="start_session_resume_notebook", + kernel_name=kernel_name, ) @mcp.tool async def start_session_continue_notebook( - notebook_path: str, experiment_name: Optional[str] = None + notebook_path: str, experiment_name: Optional[str] = None, + kernel_name: str = "python3", ) -> Dict[str, Any]: """ Start a session by continuing from an existing notebook (creates a new notebook file). @@ -338,6 +351,7 @@ async def start_session_continue_notebook( Args: notebook_path: Path to the existing notebook to continue from experiment_name: Optional custom name for the new notebook + kernel_name: Jupyter kernel spec name. Should match the notebook's language. Returns: Dictionary with: @@ -356,6 +370,7 @@ async def start_session_continue_notebook( notebook_path=notebook_path, fork_prev_notebook=True, tool_name="start_session_continue_notebook", + kernel_name=kernel_name, ) diff --git a/scribe/notebook/notebook_server.py b/scribe/notebook/notebook_server.py index 36cf625..8a97f9e 100644 --- a/scribe/notebook/notebook_server.py +++ b/scribe/notebook/notebook_server.py @@ -12,7 +12,10 @@ import nbformat from jupyter_server.serverapp import ServerApp from traitlets import Unicode, Int -from scribe.notebook._notebook_server_utils import clean_notebook_for_save +from scribe.notebook._notebook_server_utils import ( + clean_notebook_for_save, + get_notebook_metadata_for_kernel, +) from dataclasses import dataclass from . import notebook_sever_handlers as _handlers @@ -29,6 +32,7 @@ class ScribeNotebookSession: jupyter_session_id: str notebook_path: Path display_name: str + kernel_name: str = "python3" execution_count: int = 0 last_activity: Optional[datetime] = None @@ -122,6 +126,24 @@ def _setup_notebooks_directory(self) -> Path: print(f"❌ {error_msg}") raise ValueError(error_msg) from e + async def _get_kernel_language_info(self, kernel_id: str) -> dict: + """Query a running kernel for its language_info via kernel_info_request. + + Returns the full language_info dict (name, version, mimetype, etc.) + or a minimal fallback if the request times out. + """ + kernel = self.kernel_manager.get_kernel(kernel_id) + client = kernel.client() + client.start_channels() + try: + msg = await client._async_kernel_info() + return msg.get("content", {}).get("language_info", {}) + except Exception as e: + print(f"Warning: failed to query language_info from kernel {kernel_id}: {e}") + return {} + finally: + client.stop_channels() + def init_webapp(self): """Add our custom handlers to the web app.""" super().init_webapp() @@ -148,7 +170,7 @@ def init_webapp(self): self.shutdown_check_callback.start() async def start_session( - self, experiment_name=None, existing_notebook_path=None, fork_prev_notebook=True + self, experiment_name=None, existing_notebook_path=None, fork_prev_notebook=True, kernel_name="python3", ): """ Start a new scribe jupyter session -- one-to-one with a notebook and a kernel. @@ -162,6 +184,7 @@ async def start_session( NOTE: re-running cells without forking will update existing file, and may overwrite outputs e.g. with different random outputs if seeds have not been set, or with errors if incorrect env is being used. + kernel_name: Jupyter kernel spec name (e.g. "python3", "ir", "julia-1.12"). Must be installed and registered with Jupyter. """ # Ensure notebooks directory is set up if self.notebooks_path is None: @@ -225,16 +248,10 @@ async def start_session( # Create empty notebook nb = nbformat.v4.new_notebook() - nb.metadata.update( - { - "kernelspec": { - "display_name": f"Scribe: {base_name}", - "language": "python", - "name": "python3", - }, - "language_info": {"name": "python", "version": "3.11"}, - } + kernel_meta = get_notebook_metadata_for_kernel( + self.kernel_spec_manager, kernel_name ) + nb.metadata.update(kernel_meta) # Save notebook with open(nb_path, "w") as f: @@ -252,7 +269,7 @@ async def start_session( relative_path = nb_path # Create a kernel first to ensure it uses our current environment - kernel_id = await self.kernel_manager.start_kernel() + kernel_id = await self.kernel_manager.start_kernel(kernel_name=kernel_name) # Now create a session and associate it with our kernel sm = self.web_app.settings["session_manager"] @@ -271,10 +288,20 @@ async def start_session( jupyter_session_id=jupyter_session_id, notebook_path=nb_path, display_name=kernel_display_name, + kernel_name=kernel_name, last_activity=datetime.now(), ) self.sessions[session_id] = scribe_session + # Update notebook with kernel info after beginning + language_info = await self._get_kernel_language_info(kernel_id) + if language_info: + with open(nb_path, "r") as f: + nb = nbformat.read(f, as_version=nbformat.NO_CONVERT) + nb.metadata["language_info"] = language_info + with open(nb_path, "w") as f: + nbformat.write(clean_notebook_for_save(nb), f) + """STEP 3: If we have an existing notebook, execute all code cells to restore state """ restoration_results = [] if existing_notebook_path: @@ -379,6 +406,7 @@ async def start_session( "kernel_id": kernel_id, "notebook_path": str(nb_path), "kernel_display_name": kernel_display_name, + "kernel_name": kernel_name, } # Include restoration results in output if we restored from an existing notebook diff --git a/scribe/notebook/notebook_sever_handlers.py b/scribe/notebook/notebook_sever_handlers.py index cf2c160..b8fd2c3 100644 --- a/scribe/notebook/notebook_sever_handlers.py +++ b/scribe/notebook/notebook_sever_handlers.py @@ -43,6 +43,7 @@ async def post(self): fork_prev_notebook=data.get( "fork_prev_notebook", True ), # Default to True for backward compatibility + kernel_name=data.get("kernel_name", "python3"), ) # Add server URL and token to response @@ -53,7 +54,7 @@ async def post(self): result["vscode_url"] = ( f"http://localhost:{self.scribe_app.port}/?token={self.scribe_app.token}" ) - result["kernel_name"] = result.pop("kernel_display_name", "") + result["kernel_display_name"] = result.pop("kernel_display_name", "") result["status"] = "started" self.finish(json.dumps(result))