"""Message handler for JSON-RPC routing."""

from __future__ import annotations

import asyncio
import base64
import contextvars
import json
import mimetypes
import os
from .compat import get_default_system_path, get_subprocess_kwargs, kill_process_tree

import time
import uuid as uuid_mod
from datetime import datetime, timezone
from typing import Callable, Dict, Any, Optional, Awaitable, TYPE_CHECKING
import logging

from .structured_log import get_logger as _get_structured_logger, bind_trace, clear_contextvars

if TYPE_CHECKING:
    from .output_buffer import OutputBuffer

from pathlib import Path

from .kernel_backend import KernelBackend
from .execution_queue import ExecutionQueue
from .file_manager import FileManager
from .env_manager import EnvironmentManager
from .session_manager import SessionManager
from .notebook_handler import (
    read_notebook,
    write_notebook,
    write_notebook_safe,
    validate_notebook,
    create_empty_notebook,
    update_cell,
    update_cell_outputs,
    add_cell,
    delete_cell,
    CellSourceConflictError,
)

logger = logging.getLogger("skop-runner")
_rpc_log = _get_structured_logger("rpc")

# Type for notification sender callback
NotificationSender = Callable[[str, dict], Awaitable[None]]

# Maximum sizes for input validation (prevents resource exhaustion attacks)
MAX_CODE_SIZE = 1 * 1024 * 1024  # 1MB max code per execution
MAX_FILE_CONTENT_SIZE = 50 * 1024 * 1024  # 50MB max file content
MAX_NOTEBOOK_SIZE = 100 * 1024 * 1024  # 100MB max notebook (they can be large with outputs)
MAX_PATCH_SIZE = 10 * 1024 * 1024  # 10MB max patch size
MAX_CELL_SOURCE_SIZE = 1 * 1024 * 1024  # 1MB max cell source
MAX_BASH_OUTPUT = 1 * 1024 * 1024  # 1MB max per stream (stdout/stderr)
MAX_BASH_TIMEOUT = 300  # 5 minutes hard max
MAX_FIGURE_BASE64_BYTES = 5 * 1024 * 1024  # 5MB max base64 image to auto-save
MAX_CONCURRENT_BASH = 3  # Prevent fork bombs
SAFE_ENV_KEYS = {
    "PATH", "HOME", "USER", "LANG", "LC_ALL", "SHELL", "TERM",
    "VIRTUAL_ENV", "CONDA_PREFIX", "CONDA_DEFAULT_ENV", "PYTHONPATH",
    "TMPDIR", "TMP", "TEMP",
    "SSL_CERT_FILE", "SSL_CERT_DIR", "REQUESTS_CA_BUNDLE", "CURL_CA_BUNDLE",
    "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY",
    "http_proxy", "https_proxy", "no_proxy",
    # Windows-required: cmd.exe and subprocesses need these to function
    "SystemRoot", "COMSPEC", "USERPROFILE", "APPDATA", "LOCALAPPDATA",
    "HOMEDRIVE", "HOMEPATH", "USERNAME", "PATHEXT",
}

# Per-task request source — isolated across concurrent asyncio tasks via contextvars.
# Each create_task() call gets an independent copy, so concurrent handlers
# can't corrupt each other's source tag.
_request_source: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "_request_source", default=None
)

# Declarative mapping: RPC method → (notification_method, notification_params).
# Only methods listed here emit notifications, and only for agent-initiated requests.
# The lambda receives (params, result) from the handler and returns the notification tuple.
_MUTATION_NOTIFICATIONS: Dict[str, Callable[[dict, Any], tuple[str, dict]]] = {
    "notebook.updateCell": lambda params, result: (
        "notebook.cellUpdated",
        {"path": params["path"], "cell_id": params["cell_id"], "source": params["source"]},
    ),
    "notebook.addCell": lambda params, result: (
        "notebook.cellAdded",
        {"path": params.get("path"), "cell_id": result["cell_id"],
         "cell_type": params.get("cell_type"), "source": params.get("source", ""),
         "after_cell_id": params.get("after_cell_id")},
    ),
    "notebook.deleteCell": lambda params, result: (
        "notebook.cellDeleted",
        {"path": params.get("path"), "cell_id": params.get("cell_id")},
    ),
    "file.write": lambda params, result: ("file.written", {"path": params["path"]}),
    "file.edit": lambda params, result: ("file.written", {"path": params["path"]}),
    "file.create": lambda params, result: ("file.created", {"path": params["path"], "type": "file"}),
    "file.createFolder": lambda params, result: ("file.created", {"path": params["path"], "type": "directory"}),
    "file.delete": lambda params, result: ("file.deleted", {"path": params["path"]}),
    "file.rename": lambda params, result: (
        "file.renamed", {"old_path": params["old_path"], "new_path": params["new_path"]},
    ),
    "file.move": lambda params, result: (
        "file.moved", {"old_path": params["source"], "new_path": params["dest"]},
    ),
}


class MessageHandler:
    """Routes JSON-RPC messages to appropriate handlers."""

    def __init__(
        self,
        kernel_backend: KernelBackend,
        file_manager: FileManager,
        output_buffer: Optional["OutputBuffer"] = None,
        persistence_dir: Optional[str] = None,
        idle_timeout: Optional[int] = None,
    ):
        self.kernel_backend = kernel_backend
        self.file_manager = file_manager
        self.env_manager = EnvironmentManager(file_manager.get_project_root())
        self.session_manager = SessionManager(
            kernel_backend, persistence_dir=persistence_dir
        )
        self._send_notification: Optional[NotificationSender] = None
        self._output_buffer = output_buffer
        self._execution_queue = ExecutionQueue()
        self._sessions_restored = False
        self._active_bash_count = 0

        # Idle timeout for cloud runners (seconds, None = disabled)
        self._idle_timeout = idle_timeout
        self._last_activity = time.monotonic()
        self._idle_monitor_task: Optional[asyncio.Task] = None
        self._shutdown_callback: Optional[Callable] = None
        self._reconnect_callback: Optional[Callable] = None

        # Method routing map
        self.methods: Dict[str, Callable] = {
            # Kernel operations
            "kernel.start": self._handle_kernel_start,
            "kernel.stop": self._handle_kernel_stop,
            "kernel.restart": self._handle_kernel_restart,
            "kernel.interrupt": self._handle_kernel_interrupt,
            "kernel.execute": self._handle_kernel_execute,
            "kernel.executeBatch": self._handle_kernel_execute_batch,
            "kernel.inputReply": self._handle_kernel_input_reply,
            "kernel.status": self._handle_kernel_status,
            "kernel.list": self._handle_kernel_list,
            "kernel.listAvailable": self._handle_kernel_list_available,
            "kernel.installIpykernel": self._handle_kernel_install_ipykernel,
            # File operations
            "file.list": self._handle_file_list,
            "file.read": self._handle_file_read,
            "file.write": self._handle_file_write,
            "file.create": self._handle_file_create,
            "file.createFolder": self._handle_file_create_folder,
            "file.delete": self._handle_file_delete,
            "file.rename": self._handle_file_rename,
            "file.copy": self._handle_file_copy,
            "file.move": self._handle_file_move,
            "file.info": self._handle_file_info,
            "file.exists": self._handle_file_exists,
            "file.applyPatch": self._handle_file_apply_patch,
            "file.edit": self._handle_file_edit,
            "file.glob": self._handle_file_glob,
            "file.grep": self._handle_file_grep,
            "file.findNotebooks": self._handle_file_find_notebooks,
            "file.readBinary": self._handle_file_read_binary,
            # Workspace operations
            "workspace.getRoot": self._handle_workspace_get_root,
            "workspace.setRoot": self._handle_workspace_set_root,
            "workspace.browse": self._handle_workspace_browse,
            "workspace.createFolder": self._handle_workspace_create_folder,
            "workspace.home": self._handle_workspace_home,
            # Environment operations
            "env.packages": self._handle_env_packages,
            # Notebook operations (nbformat-based)
            "notebook.read": self._handle_notebook_read,
            "notebook.write": self._handle_notebook_write,
            "notebook.create": self._handle_notebook_create,
            "notebook.validate": self._handle_notebook_validate,
            "notebook.updateCell": self._handle_notebook_update_cell,
            "notebook.addCell": self._handle_notebook_add_cell,
            "notebook.deleteCell": self._handle_notebook_delete_cell,
            # Execution queue operations
            "execution.cancel": self._handle_execution_cancel,
            "execution.queue": self._handle_execution_queue,
            # State operations
            "state.snapshot": self._handle_state_snapshot,
            # System operations
            "system.bash": self._handle_system_bash,
            "system.reconnect": self._handle_system_reconnect,
            # Session operations
            "session.create": self._handle_session_create,
            "session.get": self._handle_session_get,
            "session.list": self._handle_session_list,
            "session.close": self._handle_session_close,
            "session.connectToKernel": self._handle_session_connect_to_kernel,
            "session.updatePath": self._handle_session_update_path,
        }

    def request_shutdown(self, reason: str):
        """Request graceful shutdown from any context (idle timeout, balance exhaustion, etc.)."""
        logger.info(f"Shutdown requested: {reason}")
        if self._shutdown_callback:
            self._shutdown_callback()

    def set_notification_sender(self, sender: NotificationSender):
        """Set callback for sending JSON-RPC notifications."""
        self._send_notification = sender

    async def send_notification(self, method: str, params: dict):
        """Send a JSON-RPC notification (message without id)."""
        if self._send_notification:
            await self._send_notification(method, params)

    async def _save_figure_from_outputs(
        self, cell_id: str, outputs: list[dict]
    ) -> Optional[str]:
        """Extract first image/png from outputs, save to .skop/figures/.

        Returns relative path (from project root) to saved figure, or None.
        """
        if not self.file_manager.project_root:
            return None

        for out in outputs[:10]:
            otype = out.get("output_type", "")
            if otype not in ("display_data", "execute_result"):
                continue
            data = out.get("data")
            if not isinstance(data, dict):
                continue
            img_b64 = data.get("image/png")
            if not img_b64 or not isinstance(img_b64, str):
                continue
            if len(img_b64) > MAX_FIGURE_BASE64_BYTES:
                logger.warning(
                    f"Skipping oversized figure for cell {cell_id}: "
                    f"{len(img_b64) / (1024*1024):.1f}MB"
                )
                continue

            try:
                figures_dir = self.file_manager.project_root / ".skop" / "figures"
                figures_dir.mkdir(parents=True, exist_ok=True)

                ts = int(datetime.now(timezone.utc).timestamp())
                safe_cell_id = "".join(
                    c for c in cell_id if c.isalnum() or c in "-_"
                )[:50]
                filename = f"{safe_cell_id}_{ts}.png"
                filepath = figures_dir / filename

                img_bytes = base64.b64decode(img_b64)
                filepath.write_bytes(img_bytes)

                rel_path = f".skop/figures/{filename}"
                logger.info(f"Saved figure: {rel_path}")
                return rel_path
            except Exception as e:
                logger.warning(f"Failed to save figure for cell {cell_id}: {e}")
                return None

        return None

    async def _send_context_notification(
        self,
        cell_id: Optional[str],
        code: str,
        result: dict,
        session,
    ):
        """Send context.executionCompleted notification for journal capture.

        Only fires when cell_id is present (excludes introspection tools)
        and a project_id is known from .skop/project.json.
        """
        if not cell_id or not self._send_notification:
            return
        _project_id = self.file_manager.get_project_id()
        if not _project_id:
            return
        try:
            _source = _request_source.get() or "user"

            # Auto-capture first image/png figure
            _figure_path = await self._save_figure_from_outputs(
                cell_id, result.get("outputs", [])
            )

            # Trim outputs to avoid bloating the WS transport.
            _MAX_IMG_B64 = 2 * 1024 * 1024
            _MAX_TEXT_LEN = 4000
            raw_outputs = result.get("outputs", [])[:10]
            trimmed = []
            for out in raw_outputs:
                otype = out.get("output_type", "")
                if otype == "stream":
                    text = out.get("text", "")
                    if isinstance(text, str) and len(text) > _MAX_TEXT_LEN:
                        out = {**out, "text": text[:_MAX_TEXT_LEN]}
                elif otype in ("display_data", "execute_result"):
                    data = out.get("data")
                    if isinstance(data, dict):
                        new_data = {}
                        for k, v in data.items():
                            if k.startswith("image/") and isinstance(v, str) and len(v) > _MAX_IMG_B64:
                                continue  # drop oversized images
                            if k == "text/plain" and isinstance(v, str) and len(v) > _MAX_TEXT_LEN:
                                v = v[:_MAX_TEXT_LEN]
                            new_data[k] = v
                        out = {**out, "data": new_data}
                elif otype == "error":
                    evalue = out.get("evalue", "")
                    if isinstance(evalue, str) and len(evalue) > _MAX_TEXT_LEN:
                        out = {**out, "evalue": evalue[:_MAX_TEXT_LEN]}
                trimmed.append(out)
            await self.send_notification("context.executionCompleted", {
                "project_id": _project_id,
                "cell_id": cell_id,
                "notebook_path": self.file_manager._normalize_output_path(session.notebook_path) if session else "",
                "code": code[:5000],
                "outputs": trimmed,
                "status": result.get("status", "ok"),
                "source": _source,
                "executed_at": datetime.now(timezone.utc).isoformat(),
                "figure_path": _figure_path,
            })
        except Exception as e:
            logger.warning(
                f"Failed to send context.executionCompleted for cell {cell_id}: {e}"
            )

    async def _persist_cell_outputs(
        self,
        cell_id: str,
        kernel_id: str,
        outputs: list,
        execution_count: int,
    ) -> None:
        """Persist cell outputs to the notebook file on disk.

        Fire-and-forget: errors are logged but never raised, so this
        never disrupts the execute response flow.
        """
        try:
            session = self.session_manager.get_session_for_kernel(kernel_id)
            if not session:
                return
            full_path = self.file_manager.resolve_path(session.notebook_path)
            await update_cell_outputs(
                full_path, cell_id, outputs, execution_count
            )
        except Exception as e:
            logger.warning(
                f"Failed to persist outputs for cell {cell_id}: {e}"
            )

    def _validate_id(self, msg_id) -> bool:
        """Validate JSON-RPC id field. Must be string, int, or None per spec."""
        return msg_id is None or isinstance(msg_id, (str, int))

    def _sanitize_error(self, error: Exception) -> str:
        """Sanitize error message to avoid leaking internal paths or sensitive info."""
        error_str = str(error)
        # Don't expose full file paths in errors - just the relative path or filename
        # This is a simple sanitization; complex paths may still leak
        if self.file_manager.project_root:
            root_str = str(self.file_manager.project_root)
            error_str = error_str.replace(root_str, "<project>")
        return error_str

    def start_idle_monitor(self):
        """Start background task that exits process after idle timeout."""
        if self._idle_timeout is None:
            return
        self._idle_monitor_task = asyncio.create_task(self._idle_monitor_loop())
        logger.info(f"Idle monitor started ({self._idle_timeout}s timeout)")

    async def _idle_monitor_loop(self):
        """Check for inactivity every 60s, signal shutdown if idle too long."""
        while True:
            await asyncio.sleep(60)
            elapsed = time.monotonic() - self._last_activity
            if elapsed >= self._idle_timeout:
                idle_min = int(elapsed / 60)
                logger.info(f"Idle for {idle_min}m, shutting down (cloud runner auto-stop)")
                # Send notification before exiting so browser knows why
                try:
                    await self.send_notification("runner.shutting_down", {
                        "reason": "idle_timeout",
                        "idle_minutes": idle_min,
                    })
                except Exception:
                    pass
                await asyncio.sleep(1)
                # Signal the main run loop to exit gracefully
                self.request_shutdown("idle_timeout")
                return

    async def handle_message(self, message: dict) -> dict:
        """
        Handle incoming JSON-RPC message.

        Auto-instruments all RPC handlers with structured logging:
        - rpc.start: logged on entry with method, params_keys, source
        - rpc.complete: logged on success with method, duration_ms
        - rpc.error: logged on failure with method, error_type, duration_ms

        Trace IDs are extracted from params._trace and bound to structlog
        contextvars for the duration of the request.
        """
        # Track activity for idle timeout
        self._last_activity = time.monotonic()

        msg_id = message.get("id")

        # Validate JSON-RPC id field (must be string, int, or null per spec)
        if not self._validate_id(msg_id):
            logger.warning(f"Invalid JSON-RPC id type: {type(msg_id).__name__}")
            msg_id = None  # Use null for response if id is invalid

        # Extract source tag for agent vs browser discrimination
        source = message.get("source")
        _request_source.set(source)
        t0 = time.perf_counter()
        method = "unknown"

        try:
            method = message.get("method")
            params = message.get("params") or {}

            if not method:
                return self._error_response(msg_id, -32600, "Missing method")

            if not isinstance(method, str):
                return self._error_response(msg_id, -32600, "Method must be a string")

            if method not in self.methods:
                return self._error_response(msg_id, -32601, f"Unknown method: {method}")

            # ── Trace ID + structured logging ──────────────────────────
            # Extract _trace from params (injected by browser/agent), bind to context.
            # Strip _trace before forwarding to handler — it's metadata, not a param.
            clear_contextvars()
            trace_id = params.pop("_trace", None) or str(uuid_mod.uuid4())
            bind_trace(trace_id)

            _rpc_log.info(
                "rpc.start",
                method=method,
                params_keys=list(params.keys()),
                source=source or "unknown",
            )
            # ──────────────────────────────────────────────────────────

            handler = self.methods[method]
            result = await handler(params)

            # Log success
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.info("rpc.complete", method=method, duration_ms=duration_ms)

            # Emit mutation notification for agent-initiated changes
            if _request_source.get() == "agent" and method in _MUTATION_NOTIFICATIONS:
                try:
                    notif_method, notif_params = _MUTATION_NOTIFICATIONS[method](params, result)
                    await self.send_notification(notif_method, notif_params)
                    _rpc_log.debug("mutation.notification_sent", notif_method=notif_method)
                except Exception as notif_err:
                    _rpc_log.warning(
                        "mutation.notification_failed",
                        notif_method=method,
                        error=str(notif_err),
                    )

            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": result,
            }

        except CellSourceConflictError as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.info("rpc.error", method=method, error_type="CellSourceConflictError", duration_ms=duration_ms)
            return self._error_response(
                msg_id, -32009, "Cell source changed",
                data={"current_source": e.current_source},
            )

        except KeyError as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.warning("rpc.error", method=method, error_type="KeyError", error=str(e), duration_ms=duration_ms)
            return self._error_response(msg_id, -32602, f"Missing parameter: {e}")

        except FileNotFoundError as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.warning("rpc.error", method=method, error_type="FileNotFoundError", duration_ms=duration_ms)
            return self._error_response(msg_id, -32001, self._sanitize_error(e))

        except FileExistsError as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.warning("rpc.error", method=method, error_type="FileExistsError", duration_ms=duration_ms)
            return self._error_response(msg_id, -32002, self._sanitize_error(e))

        except ValueError as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.warning("rpc.error", method=method, error_type="ValueError", error=str(e), duration_ms=duration_ms)
            return self._error_response(msg_id, -32003, self._sanitize_error(e))

        except Exception as e:
            duration_ms = round((time.perf_counter() - t0) * 1000, 2)
            _rpc_log.error("rpc.error", method=method, error_type=type(e).__name__, duration_ms=duration_ms, exc_info=True)
            return self._error_response(msg_id, -32000, self._sanitize_error(e))

    def _error_response(
        self,
        msg_id: Optional[str],
        code: int,
        message: str,
        data: Optional[dict] = None,
    ) -> dict:
        """Create JSON-RPC error response."""
        error: Dict[str, Any] = {
            "code": code,
            "message": message,
        }
        if data is not None:
            error["data"] = data
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "error": error,
        }

    # ========== Kernel handlers ==========

    async def _handle_kernel_start(self, params: dict) -> dict:
        """
        Start a new kernel.

        Params:
            kernel_name: Jupyter kernel spec name (default: python3)
            python_path: Optional path to Python executable for custom venvs
            env_name: Optional display name for the environment
        """
        kernel_name = params.get("kernel_name", "python3")
        python_path = params.get("python_path")
        env_name = params.get("env_name")

        # Cloud runners: auto-select persistent venv if no explicit python_path
        if not python_path and os.getenv("SKOP_CLOUD_RUNNER") == "true":
            venv_python = Path("/workspace/.venv/bin/python")
            if venv_python.exists():
                python_path = str(venv_python)
                env_name = env_name or ".venv"

        cwd = self.file_manager.get_project_root()
        kernel_id = await self.kernel_backend.start_kernel(
            kernel_name=kernel_name,
            python_path=python_path,
            env_name=env_name,
            cwd=cwd,
        )

        # Get kernel info to return the actual python_path and env_name used
        kernel_info = self.kernel_backend.get_kernel_info(kernel_id)
        actual_python_path = kernel_info.python_path if kernel_info else python_path
        actual_env_name = kernel_info.env_name if kernel_info else env_name

        return {
            "kernel_id": kernel_id,
            "python_path": actual_python_path,
            "env_name": actual_env_name,
        }

    async def _handle_kernel_stop(self, params: dict) -> dict:
        """Stop a kernel."""
        await self.kernel_backend.stop_kernel(params["kernel_id"])
        return {"success": True}

    async def _handle_kernel_restart(self, params: dict) -> dict:
        """Restart a kernel."""
        await self.kernel_backend.restart_kernel(params["kernel_id"])
        return {"success": True}

    async def _handle_kernel_interrupt(self, params: dict) -> dict:
        """Interrupt a kernel."""
        await self.kernel_backend.interrupt_kernel(params["kernel_id"])
        return {"success": True}

    async def _handle_kernel_input_reply(self, params: dict) -> dict:
        """Send an input reply to a kernel waiting for stdin."""
        kernel_id = params["kernel_id"]
        value = params.get("value", "")
        await self.kernel_backend.send_input_reply(kernel_id, value)
        return {"success": True}

    async def _handle_kernel_execute(self, params: dict) -> dict:
        """
        Execute code in a kernel with streaming output support.

        Routes through the per-kernel execution queue to ensure sequential
        execution within each kernel.

        Params:
            kernel_id: Kernel to execute in
            code: Code to execute
            cell_id: Cell ID for streaming notifications (required for streaming)
            request_id: Unique execution request ID for correlating outputs

        Returns:
            outputs: List of output messages (also sent as streaming notifications)
            execution_count: The execution count for this cell
            status: "ok" or "error"
            request_id: Echo of the request_id for verification

        Streaming:
            If cell_id is provided and notification sender is configured,
            sends kernel.output notifications as outputs arrive:
            {cell_id, kernel_id, session_id, request_id, output, index}
        """
        kernel_id = params["kernel_id"]
        code = params["code"]
        cell_id = params.get("cell_id")
        request_id = params.get("request_id")

        # Security: Validate code size to prevent resource exhaustion
        if len(code) > MAX_CODE_SIZE:
            raise ValueError(
                f"Code too large: {len(code) / 1024:.1f}KB "
                f"(max {MAX_CODE_SIZE / 1024:.0f}KB)"
            )

        # Look up session for this kernel (for routing outputs to correct notebook)
        session = self.session_manager.get_session_for_kernel(kernel_id)
        session_id = session.session_id if session else None

        async def execute_fn():
            # Send execution_start notification (replaces frontend onStart callback)
            if cell_id and self._send_notification:
                try:
                    await self.send_notification("kernel.output", {
                        "cell_id": cell_id,
                        "kernel_id": kernel_id,
                        "session_id": session_id,
                        "request_id": request_id,
                        "output": {"output_type": "execution_start"},
                        "index": -1,
                    })
                except Exception:
                    logger.warning(f"Failed to send execution_start for cell {cell_id}")
            try:
                result = await self._execute_cell_streaming(
                    kernel_id, code, cell_id, request_id, session_id
                )
            except Exception as e:
                logger.error(f"Cell execution failed for {cell_id}: {e}")
                # Send execution_end with error so browser finalizes the cell
                if cell_id and self._send_notification:
                    try:
                        await self.send_notification("kernel.output", {
                            "cell_id": cell_id,
                            "kernel_id": kernel_id,
                            "session_id": session_id,
                            "request_id": request_id,
                            "output": {
                                "output_type": "execution_end",
                                "execution_count": 0,
                                "status": "error",
                            },
                            "index": -1,
                        })
                    except Exception:
                        logger.warning(f"Failed to send execution_end for cell {cell_id}")
                raise
            result["request_id"] = request_id
            # Send execution_end so browser can finalize (especially for
            # agent-initiated runs where the RPC response doesn't reach the UI)
            if cell_id and self._send_notification:
                try:
                    await self.send_notification("kernel.output", {
                        "cell_id": cell_id,
                        "kernel_id": kernel_id,
                        "session_id": session_id,
                        "request_id": request_id,
                        "output": {
                            "output_type": "execution_end",
                            "execution_count": result.get("execution_count", 0),
                            "status": result.get("status", "ok"),
                        },
                        "index": -1,
                    })
                except Exception:
                    logger.warning(f"Failed to send execution_end for cell {cell_id}")

            # Send context.executionCompleted for journal capture
            await self._send_context_notification(
                cell_id, code, result, session
            )

            # Persist outputs to disk so they survive browser tab switches
            if cell_id:
                await self._persist_cell_outputs(
                    cell_id, kernel_id,
                    result["outputs"], result.get("execution_count", 0),
                )

            return result

        return await self._execution_queue.submit(
            kernel_id, [cell_id] if cell_id else [], execute_fn
        )

    async def _execute_cell_streaming(
        self,
        kernel_id: str,
        code: str,
        cell_id: Optional[str],
        request_id: Optional[str],
        session_id: Optional[str],
    ) -> dict:
        """
        Execute code in a kernel with streaming output notifications.

        Sends kernel.output notifications as outputs arrive, and returns
        the final result with collected outputs.

        Returns:
            {"outputs": [...], "execution_count": int, "status": "ok"|"error"}
        """
        outputs = []
        status = "ok"
        output_index = 0
        pending_clear = False

        async for output in self.kernel_backend.execute_code(kernel_id, code, timeout=0):
            output_type = output.get("output_type")

            # Send streaming notification if cell_id is provided
            if cell_id and self._send_notification:
                await self.send_notification("kernel.output", {
                    "cell_id": cell_id,
                    "kernel_id": kernel_id,
                    "session_id": session_id,
                    "request_id": request_id,
                    "output": output,
                    "index": output_index,
                })
                output_index += 1

            # input_request is transient — stream to frontend but don't persist
            if output_type == "input_request":
                continue

            if output_type == "clear_output":
                if output.get("wait"):
                    pending_clear = True
                else:
                    outputs = []
                    pending_clear = False
                continue

            if pending_clear:
                outputs = []
                pending_clear = False

            if output_type == "update_display_data":
                display_id = output.get("display_id")
                if display_id:
                    for i, existing in enumerate(outputs):
                        if existing.get("display_id") == display_id:
                            outputs[i] = {
                                **existing,
                                "data": output.get("data", {}),
                                "metadata": output.get("metadata", {}),
                            }
                            break
                continue

            outputs.append(output)

            if output_type == "error":
                status = "error"

        execution_count = self.kernel_backend.get_execution_count(kernel_id)

        # Strip display_id from outputs — runtime-only field not valid in nbformat
        for output in outputs:
            output.pop("display_id", None)

        return {
            "outputs": outputs,
            "execution_count": execution_count,
            "status": status,
        }

    async def _handle_kernel_execute_batch(self, params: dict) -> dict:
        """
        Execute multiple cells sequentially in a kernel with streaming output.

        The entire batch is submitted as a single queue slot, preventing
        interleaving with standalone kernel.execute calls.

        Params:
            kernel_id: Kernel to execute in
            cells: List of {cell_id, code, request_id} dicts
            stop_on_error: Stop executing remaining cells on first error (default: True)

        Returns:
            results: List of {cell_id, outputs, execution_count, status, request_id}
        """
        kernel_id = params["kernel_id"]
        cells = params["cells"]
        stop_on_error = params.get("stop_on_error", True)

        # Look up session for this kernel
        session = self.session_manager.get_session_for_kernel(kernel_id)
        session_id = session.session_id if session else None
        cell_ids = [c["cell_id"] for c in cells]

        async def execute_fn():
            results = []

            for cell in cells:
                cell_id = cell["cell_id"]
                code = cell["code"]
                request_id = cell.get("request_id")

                try:
                    # Security: Validate code size
                    if len(code) > MAX_CODE_SIZE:
                        raise ValueError(
                            f"Code too large: {len(code) / 1024:.1f}KB "
                            f"(max {MAX_CODE_SIZE / 1024:.0f}KB)"
                        )

                    # Send execution_start notification so frontend transitions queued -> running
                    if self._send_notification:
                        await self.send_notification("kernel.output", {
                            "cell_id": cell_id,
                            "kernel_id": kernel_id,
                            "session_id": session_id,
                            "request_id": request_id,
                            "output": {"output_type": "execution_start"},
                            "index": -1,
                        })

                    result = await self._execute_cell_streaming(
                        kernel_id, code, cell_id, request_id, session_id
                    )
                except Exception as e:
                    # Kernel crash, notification failure, or validation error —
                    # return partial results so the frontend can finalize cells
                    # that already completed
                    logger.error(f"Cell execution failed for {cell_id}: {e}")
                    # Notify frontend so cell finalizes immediately
                    if self._send_notification:
                        try:
                            await self.send_notification("kernel.output", {
                                "cell_id": cell_id,
                                "kernel_id": kernel_id,
                                "session_id": session_id,
                                "request_id": request_id,
                                "output": {
                                    "output_type": "execution_end",
                                    "execution_count": 0,
                                    "status": "error",
                                },
                                "index": -1,
                            })
                        except Exception:
                            pass
                    results.append({
                        "cell_id": cell_id,
                        "outputs": [],
                        "execution_count": 0,
                        "status": "error",
                        "request_id": request_id,
                    })
                    break

                cell_result = {
                    "cell_id": cell_id,
                    "outputs": result["outputs"],
                    "execution_count": result["execution_count"],
                    "status": result["status"],
                    "request_id": request_id,
                }
                results.append(cell_result)

                # Notify frontend so cell finalizes immediately
                if self._send_notification:
                    try:
                        await self.send_notification("kernel.output", {
                            "cell_id": cell_id,
                            "kernel_id": kernel_id,
                            "session_id": session_id,
                            "request_id": request_id,
                            "output": {
                                "output_type": "execution_end",
                                "execution_count": result["execution_count"],
                                "status": result["status"],
                            },
                            "index": -1,
                        })
                    except Exception:
                        logger.warning(f"Failed to send execution_end for cell {cell_id}")

                # Context notification for journal capture
                await self._send_context_notification(
                    cell_id, code, result, session
                )

                # Persist outputs to disk so they survive browser tab switches
                await self._persist_cell_outputs(
                    cell_id, kernel_id,
                    result["outputs"], result.get("execution_count", 0),
                )

                if stop_on_error and result["status"] == "error":
                    break

            return {"results": results}

        return await self._execution_queue.submit(kernel_id, cell_ids, execute_fn)

    async def _handle_kernel_status(self, params: dict) -> dict:
        """Get kernel status."""
        return await self.kernel_backend.get_kernel_status(params["kernel_id"])

    # ========== Execution queue handlers ==========

    async def _handle_execution_cancel(self, params: dict) -> dict:
        """
        Cancel all queued executions for a kernel and interrupt the running one.

        Combines queue-clear + kernel-interrupt into one atomic operation.

        Params:
            kernel_id: Kernel whose queue to cancel
        """
        kernel_id = params["kernel_id"]
        result = await self._execution_queue.cancel_queue(
            kernel_id, interrupt_fn=self.kernel_backend.interrupt_kernel
        )
        return {"success": True, **result}

    async def _handle_execution_queue(self, params: dict) -> dict:
        """
        Get execution queue state for one or all kernels.

        Params:
            kernel_id: Optional kernel ID (returns all if omitted)
        """
        kernel_id = params.get("kernel_id")
        return self._execution_queue.get_queue_state(kernel_id)

    async def _handle_kernel_list(self, params: dict) -> dict:
        """List all running kernels."""
        return {"kernels": self.kernel_backend.list_kernels()}

    async def _handle_kernel_list_available(self, params: dict) -> dict:
        """
        List all available environments/kernels that can be started.

        Uses VS Code-style environment-first approach:
        - Shows detected Python environments (venv, conda, pyenv, poetry, pipenv, system)
        - Display names include Python version (e.g., "Python 3.11.5 (.venv)")
        - Hides Python kernelspecs (they duplicate detected environments)
        - Shows non-Python kernelspecs (R, Julia, etc.)

        Returns:
            Dict with 'environments' list containing:
            - name: Environment identifier
            - display_name: Human-readable name with Python version
            - type: "venv", "conda", "pyenv", "poetry", "pipenv", "system",
                    or "kernelspec:<language>" for non-Python
            - python_path: Path to Python executable (if known)
            - python_version: Version string (if known)
            - has_ipykernel: Whether ipykernel is installed
            - kernel_name: None for Python envs (uses ephemeral registration),
                          or kernelspec name for non-Python
            - language: "python" or the kernel language
        """
        # Try to get kernelspecs from backend (non-blocking if not running)
        # This is only used to discover non-Python kernels (R, Julia, etc.)
        # get_kernelspecs() is an optional backend capability (not part of protocol)
        kernelspecs = None
        if hasattr(self.kernel_backend, 'get_kernelspecs'):
            try:
                kernelspecs = await self.kernel_backend.get_kernelspecs()
            except Exception as e:
                logger.debug(f"Could not get kernelspecs from backend: {e}")

        # Use environment-first approach (VS Code style)
        force_refresh = params.get("refresh", False)
        environments = self.env_manager.get_available_environments(
            kernelspecs=kernelspecs,
            force_refresh=force_refresh,
        )

        return {"environments": environments}

    async def _handle_kernel_install_ipykernel(self, params: dict) -> dict:
        """
        Install ipykernel in a Python environment.

        Params:
            python_path: Path to the Python executable
        """
        python_path = params["python_path"]
        return await self.env_manager.install_ipykernel(python_path)

    # ========== File handlers ==========

    async def _handle_file_list(self, params: dict) -> dict:
        """List directory contents."""
        path = params.get("path", "")
        return {"items": await self.file_manager.list_directory(path)}

    async def _handle_file_read(self, params: dict) -> dict:
        """Read file contents."""
        content = await self.file_manager.read_file(params["path"])
        return {"content": content}

    async def _handle_file_read_binary(self, params: dict) -> dict:
        """Read file as base64-encoded binary with MIME type detection."""
        path = params["path"]
        max_size = min(
            params.get("maxSize", MAX_FILE_CONTENT_SIZE), MAX_FILE_CONTENT_SIZE
        )

        # Resolve and check size BEFORE reading into memory.
        # read_file_bytes() also calls resolve_path() internally,
        # but we need the resolved path here to check the caller's
        # maxSize (which may be lower than the 50 MB hard limit).
        file_path = self.file_manager.resolve_path(path)
        if not file_path.is_file():
            raise FileNotFoundError(f"File not found: {path}")
        file_size = file_path.stat().st_size
        if file_size > max_size:
            raise ValueError(
                f"File too large: {file_size / (1024*1024):.1f}MB "
                f"(max {max_size / (1024*1024):.0f}MB)"
            )

        data = await self.file_manager.read_file_bytes(path)
        mime_type = mimetypes.guess_type(path)[0] or "application/octet-stream"
        encoded = base64.b64encode(data).decode("ascii")

        return {"content": encoded, "mimeType": mime_type, "size": file_size}

    async def _handle_file_write(self, params: dict) -> dict:
        """Write file contents."""
        path = params["path"]
        content = params["content"]

        # Security: Validate content size to prevent resource exhaustion
        if len(content) > MAX_FILE_CONTENT_SIZE:
            raise ValueError(
                f"Content too large: {len(content) / (1024*1024):.1f}MB "
                f"(max {MAX_FILE_CONTENT_SIZE / (1024*1024):.0f}MB)"
            )

        await self.file_manager.write_file(path, content)
        return {"success": True}

    async def _handle_file_create(self, params: dict) -> dict:
        """Create empty file."""
        path = params["path"]
        await self.file_manager.create_file(path)
        return {"success": True}

    async def _handle_file_create_folder(self, params: dict) -> dict:
        """Create folder."""
        path = params["path"]
        await self.file_manager.create_folder(path)
        return {"success": True}

    async def _handle_file_delete(self, params: dict) -> dict:
        """Delete file or folder."""
        path = params["path"]
        await self.file_manager.delete(path)
        return {"success": True}

    async def _handle_file_rename(self, params: dict) -> dict:
        """Rename/move file or folder."""
        old_path = params["old_path"]
        new_path = params["new_path"]
        await self.file_manager.rename(old_path, new_path)
        return {"success": True}

    async def _handle_file_copy(self, params: dict) -> dict:
        """Copy file or folder."""
        await self.file_manager.copy_item(params["source"], params["dest"])
        return {"success": True}

    async def _handle_file_move(self, params: dict) -> dict:
        """Move file or folder."""
        source = params["source"]
        dest = params["dest"]
        await self.file_manager.move_item(source, dest)
        return {"success": True}

    async def _handle_file_info(self, params: dict) -> dict:
        """Get file info."""
        return {"info": self.file_manager.get_file_info(params["path"])}

    async def _handle_file_exists(self, params: dict) -> dict:
        """Check if file exists."""
        return {"exists": self.file_manager.exists(params["path"])}

    async def _handle_file_apply_patch(self, params: dict) -> dict:
        """Apply unified diff patch to a file."""
        patch = params["patch"]

        # Security: Validate patch size to prevent resource exhaustion
        if len(patch) > MAX_PATCH_SIZE:
            raise ValueError(
                f"Patch too large: {len(patch) / (1024*1024):.1f}MB "
                f"(max {MAX_PATCH_SIZE / (1024*1024):.0f}MB)"
            )

        result = await self.file_manager.apply_patch(params["path"], patch)
        return result

    async def _handle_file_edit(self, params: dict) -> dict:
        """
        Perform a targeted find-and-replace edit on a file.

        Params:
            path: Relative path to the file
            old_string: The exact text to find
            new_string: The replacement text
            replace_all: Replace all occurrences (default: False, requires unique match)

        Returns:
            success: True if edit succeeded
            replacements: Number of replacements made
        """
        path = params.get("path")
        old_string = params.get("old_string")
        new_string = params.get("new_string")
        replace_all = params.get("replace_all", False)

        if not path or not isinstance(path, str):
            raise ValueError("path is required and must be a string")
        if not old_string or not isinstance(old_string, str):
            raise ValueError("old_string is required and must be a non-empty string")
        if new_string is None or not isinstance(new_string, str):
            raise ValueError("new_string is required and must be a string")
        if old_string == new_string:
            raise ValueError("old_string and new_string must be different")

        return await self.file_manager.edit_file(path, old_string, new_string, replace_all)

    async def _handle_file_glob(self, params: dict) -> dict:
        """
        Find files matching a glob pattern.

        Params:
            pattern: Glob pattern (e.g., "**/*.py", "src/**/*.ts")
            path: Optional subdirectory to search in (relative to project root)

        Returns:
            files: List of matching file paths (relative to project root)
        """
        pattern = params.get("pattern")
        path = params.get("path", "")

        if not pattern or not isinstance(pattern, str):
            raise ValueError("pattern is required and must be a string")

        files = await self.file_manager.glob_files(pattern, path)
        return {"files": files}

    async def _handle_file_grep(self, params: dict) -> dict:
        """
        Search file contents with regex.

        Params:
            pattern: Regex pattern to search for
            path: Optional subdirectory to search in (relative to project root)
            include: Optional glob pattern to filter files (e.g., "*.py")

        Returns:
            matches: List of {file, line, column, match, context}
        """
        pattern = params.get("pattern")
        path = params.get("path", "")
        include = params.get("include", "")

        if not pattern or not isinstance(pattern, str):
            raise ValueError("pattern is required and must be a string")

        matches = await self.file_manager.grep_files(pattern, path, include)
        return {"matches": matches}

    async def _handle_file_find_notebooks(self, params: dict) -> dict:
        """Recursively find all .ipynb files in the project."""
        notebooks = await self.file_manager.find_notebooks()
        return {"notebooks": notebooks}

    # ========== Workspace handlers ==========

    async def _handle_workspace_get_root(self, params: dict) -> dict:
        """Get workspace root path."""
        root = self.file_manager.get_project_root()
        return {"root": root, "hasRoot": root is not None}

    async def _handle_workspace_set_root(self, params: dict) -> dict:
        """
        Set the workspace root path.

        Workspace state (FileManager/EnvironmentManager) and kernel backend
        are loosely coupled:
        1. FileManager/EnvironmentManager are set first - always succeeds
        2. Backend root is set best-effort - failures are logged but don't
           block workspace operations

        This ensures users can browse files and inspect notebooks even if
        the kernel backend isn't available or fails to start.
        """
        # Set file/env roots first - this always succeeds and enables
        # basic workspace operations (file browsing, notebook reading, etc.)
        self.file_manager.set_root(params["path"])
        resolved_root = self.file_manager.get_project_root()
        self.env_manager.set_project_root(resolved_root)

        # Write .skop/project.json if project_id provided
        project_id = params.get("project_id")
        if project_id and isinstance(project_id, str):
            try:
                self.file_manager.write_skop_project_id(project_id)
            except Exception as e:
                logger.warning(f"Failed to write .skop/project.json: {e}")

        # Optionally shut down existing kernels when switching projects
        if params.get("shutdownKernels", True):
            await self.kernel_backend.shutdown_all()

        # Set backend root directory best-effort
        # If this fails, workspace operations still work - just no kernel execution
        server_error = None
        try:
            await self.kernel_backend.set_root_directory(resolved_root)
        except Exception as e:
            # Log error but don't fail the RPC - workspace is still usable
            server_error = str(e)
            logger.warning(
                f"Backend failed to set root directory (workspace still usable): {e}"
            )

        # Restore persisted sessions after first server connection.
        # Must happen after set_root_directory which reconnects to surviving kernels
        # so that get_kernel_connection_info can validate them.
        if not self._sessions_restored and server_error is None:
            self._sessions_restored = True
            restored = await self.session_manager.load_from_disk()
            if restored:
                logger.info(f"Restored {restored} session(s) from previous run")

        return {
            "success": True,
            "root": resolved_root,
            "serverError": server_error,  # Frontend can show warning if set
        }

    async def _handle_workspace_browse(self, params: dict) -> dict:
        """
        Browse an absolute path for project root selection.

        This does NOT require project root to be set - it's for selecting
        a project folder before any project is opened.

        Params:
            path: Absolute path to browse (defaults to home directory)
        """
        path = params.get("path") or self.file_manager.get_home_directory()
        entries = self.file_manager.browse_absolute(path)
        return {"entries": entries, "path": self.file_manager._normalize_output_path(path)}

    async def _handle_workspace_create_folder(self, params: dict) -> dict:
        """
        Create a folder at an absolute path (for project root creation).

        Params:
            path: Absolute path of the folder to create
        """
        path = params.get("path")
        if not path:
            raise ValueError("path is required")
        resolved = self.file_manager.create_folder_absolute(path)
        return {"success": True, "path": resolved}

    async def _handle_workspace_home(self, params: dict) -> dict:
        """Get the user's home directory path."""
        return {"home": self.file_manager.get_home_directory()}

    # ========== Environment handlers ==========

    async def _handle_env_packages(self, params: dict) -> dict:
        """
        List packages installed in an environment.

        Params:
            env_path: Path to environment directory (venv/conda style)
            python_path: Direct path to Python executable

        At least one of env_path or python_path must be provided.
        python_path is preferred when both are given.
        """
        env_path = params.get("env_path")
        python_path = params.get("python_path")

        if not env_path and not python_path:
            raise ValueError("Either env_path or python_path must be provided")

        packages = self.env_manager.get_installed_packages(
            env_path=env_path,
            python_path=python_path,
        )
        return {"packages": packages}

    # ========== Notebook handlers ==========

    async def _handle_notebook_read(self, params: dict) -> dict:
        """
        Read notebook with nbformat validation and version upgrade.

        Params:
            path: Relative path to .ipynb file

        Returns:
            notebook: Validated notebook dict
        """
        path = params["path"]
        full_path = self.file_manager.resolve_path(path)
        notebook = read_notebook(full_path)
        return {"notebook": notebook}

    async def _handle_notebook_write(self, params: dict) -> dict:
        """
        Write notebook with nbformat validation.

        Params:
            path: Relative path to .ipynb file
            notebook: Notebook dict to write

        Returns:
            success: True if written successfully
        """
        path = params["path"]
        notebook = params["notebook"]

        # Security: Validate notebook size to prevent resource exhaustion
        notebook_size = len(json.dumps(notebook))
        if notebook_size > MAX_NOTEBOOK_SIZE:
            raise ValueError(
                f"Notebook too large: {notebook_size / (1024*1024):.1f}MB "
                f"(max {MAX_NOTEBOOK_SIZE / (1024*1024):.0f}MB)"
            )

        # Validate first
        errors = validate_notebook(notebook)
        if errors:
            raise ValueError(f"Invalid notebook: {errors[0]}")

        full_path = self.file_manager.resolve_path(path)
        await write_notebook_safe(full_path, notebook)
        return {"success": True}

    async def _handle_notebook_create(self, params: dict) -> dict:
        """
        Create a new empty notebook.

        Params:
            path: Relative path for new .ipynb file

        Returns:
            notebook: The created notebook dict
        """
        path = params["path"]
        full_path = self.file_manager.resolve_path(path)

        if full_path.exists():
            raise FileExistsError(f"File already exists: {path}")

        notebook = create_empty_notebook()
        write_notebook(full_path, notebook)
        return {"notebook": notebook}

    async def _handle_notebook_validate(self, params: dict) -> dict:
        """
        Validate notebook structure.

        Params:
            notebook: Notebook dict to validate

        Returns:
            valid: True if valid
            errors: List of validation error messages
        """
        notebook = params["notebook"]
        errors = validate_notebook(notebook)
        return {"valid": len(errors) == 0, "errors": errors}

    async def _handle_notebook_update_cell(self, params: dict) -> dict:
        """
        Update a specific cell's source code in a notebook.

        Safer than file.applyPatch for notebooks because it:
        1. Parses the notebook with nbformat
        2. Updates only the specified cell
        3. Validates and writes back with nbformat

        Params:
            path: Relative path to notebook file
            cell_id: ID of the cell to update
            source: New source code for the cell
            expected_source: Optional optimistic concurrency check — if provided,
                the cell's current source must match or a -32009 error is returned.

        Returns:
            success: True if update succeeded
        """
        path = params.get("path")
        cell_id = params.get("cell_id")
        source = params.get("source")
        expected_source = params.get("expected_source")

        if not path or not isinstance(path, str):
            raise ValueError("path is required and must be a string")
        if not cell_id or not isinstance(cell_id, str):
            raise ValueError("cell_id is required and must be a string")
        if source is None or not isinstance(source, str):
            raise ValueError("source is required and must be a string")
        if expected_source is not None and not isinstance(expected_source, str):
            raise ValueError("expected_source must be a string")

        # Security: Validate source size to prevent resource exhaustion
        if len(source) > MAX_CELL_SOURCE_SIZE:
            raise ValueError(
                f"Cell source too large: {len(source) / 1024:.1f}KB "
                f"(max {MAX_CELL_SOURCE_SIZE / 1024:.0f}KB)"
            )

        full_path = self.file_manager.resolve_path(path)
        cell_type = await update_cell(full_path, cell_id, source, expected_source=expected_source)
        return {"success": True, "cell_type": cell_type}

    async def _handle_notebook_add_cell(self, params: dict) -> dict:
        """
        Add a new cell to a notebook.

        Params:
            path: Relative path to notebook file
            cell_type: Type of cell ("code" or "markdown")
            source: Initial source code for the cell (optional, default "")
            after_cell_id: Insert after this cell ID (optional, default appends to end)

        Returns:
            cell_id: The ID of the newly created cell
        """
        path = params.get("path")
        cell_type = params.get("cell_type")
        source = params.get("source", "")
        after_cell_id = params.get("after_cell_id")

        if not path or not isinstance(path, str):
            raise ValueError("path is required and must be a string")
        if cell_type not in ("code", "markdown"):
            raise ValueError("cell_type must be 'code' or 'markdown'")
        if not isinstance(source, str):
            raise ValueError("source must be a string")
        if after_cell_id is not None and not isinstance(after_cell_id, str):
            raise ValueError("after_cell_id must be a string or null")

        # Security: Validate source size to prevent resource exhaustion
        if len(source) > MAX_CELL_SOURCE_SIZE:
            raise ValueError(
                f"Cell source too large: {len(source) / 1024:.1f}KB "
                f"(max {MAX_CELL_SOURCE_SIZE / 1024:.0f}KB)"
            )

        full_path = self.file_manager.resolve_path(path)
        cell_id = await add_cell(full_path, cell_type, source, after_cell_id)
        return {"cell_id": cell_id}

    async def _handle_notebook_delete_cell(self, params: dict) -> dict:
        """
        Delete a cell from a notebook.

        Params:
            path: Relative path to notebook file
            cell_id: ID of the cell to delete

        Returns:
            success: True if deletion succeeded
        """
        path = params.get("path")
        cell_id = params.get("cell_id")

        if not path or not isinstance(path, str):
            raise ValueError("path is required and must be a string")
        if not cell_id or not isinstance(cell_id, str):
            raise ValueError("cell_id is required and must be a string")

        full_path = self.file_manager.resolve_path(path)
        await delete_cell(full_path, cell_id)
        return {"success": True}

    # ========== State handlers ==========

    async def _handle_state_snapshot(self, params: dict) -> dict:
        """
        Return complete runner state for frontend reconstruction.

        Called on frontend reconnect to rebuild all state from the backend.
        Returns sessions with live kernel status, running kernels, and project root.

        Returns:
            sessions: List of session dicts with execution_state
            kernels: List of running kernel dicts
            project_root: Current workspace root path (or null)
        """
        sessions = await self.session_manager.list_sessions_with_status()
        kernels = self.kernel_backend.list_kernels()
        project_root = self.file_manager.get_project_root()
        buffered_outputs = self._output_buffer.drain() if self._output_buffer else []

        return {
            "sessions": sessions,
            "kernels": kernels,
            "project_root": project_root,
            "buffered_outputs": buffered_outputs,
            "execution_queues": self._execution_queue.get_queue_state(),
        }

    # ========== System handlers ==========

    async def _handle_system_reconnect(self, params: dict) -> dict:
        """Force WebSocket reconnection for pool machine reassignment.
        Also enables idle monitor — pool machines run without timeout,
        but once assigned to a user they should auto-shutdown on idle."""
        reason = params.get("reason", "unknown")
        logger.info(f"system.reconnect requested: {reason}")

        # Enable idle timeout after pool reassignment (pool machines start without it)
        if reason == "pool_reassignment" and self._idle_timeout is None:
            self._idle_timeout = 2592000  # 30 days
            self.start_idle_monitor()
            logger.info("Idle monitor enabled after pool reassignment (30 days)")

        if self._reconnect_callback:
            async def _deferred():
                await asyncio.sleep(0.2)
                self._reconnect_callback()
            asyncio.create_task(_deferred())

        return {"status": "reconnecting", "reason": reason}

    def _resolve_project_python_bin(self) -> Optional[str]:
        """Resolve the bin/ directory for the project's active Python environment.

        Checks active kernel sessions first (already running), then falls back
        to environment detection (venv, poetry, pipenv, conda).  Returns None
        when no non-system environment is found.
        """
        # 1. Active kernel sessions — O(1), already in memory
        for s in self.session_manager.list_sessions():
            pp = s.get("python_path")
            if pp:
                bin_dir = str(Path(pp).parent)
                if Path(bin_dir).is_dir():
                    return bin_dir

        # 2. Detected environments (cached, 60s TTL)
        _VENV_TYPES = {"venv", "poetry", "pipenv", "conda"}
        for env in self.env_manager.detect_environments():
            if env.get("type") in _VENV_TYPES and env.get("python_path"):
                bin_dir = str(Path(env["python_path"]).parent)
                if Path(bin_dir).is_dir():
                    return bin_dir

        return None

    async def _handle_system_bash(self, params: dict) -> dict:
        """
        Execute a shell command in the project directory.

        Params:
            command: Shell command string to execute
            timeout: Max execution time in seconds (default: 30, max: 300)

        Returns:
            stdout: Command stdout (truncated to 1MB)
            stderr: Command stderr (truncated to 1MB)
            exit_code: Process return code (-1 on timeout)
            truncated: Whether output was truncated
        """
        command = params.get("command")
        timeout = max(1, min(params.get("timeout", 120), MAX_BASH_TIMEOUT))

        if not isinstance(command, str) or not command.strip():
            raise ValueError("command must be a non-empty string")

        if self._active_bash_count >= MAX_CONCURRENT_BASH:
            raise RuntimeError(
                f"Too many concurrent bash commands (max {MAX_CONCURRENT_BASH})"
            )

        project_root = self.file_manager.project_root
        if not project_root:
            raise ValueError("No project root set")

        self._active_bash_count += 1
        # External chunk lists so partial output survives task cancellation
        stdout_chunks: list[bytes] = []
        stderr_chunks: list[bytes] = []
        try:
            clean_env = {k: v for k, v in os.environ.items() if k in SAFE_ENV_KEYS}
            if "PATH" not in clean_env:
                clean_env["PATH"] = get_default_system_path()

            # Activate the project's Python environment so pip/python resolve
            # to the correct venv instead of system Python (or not at all).
            venv_bin = self._resolve_project_python_bin()
            if venv_bin:
                clean_env["PATH"] = f"{venv_bin}{os.pathsep}{clean_env['PATH']}"
                clean_env["VIRTUAL_ENV"] = str(Path(venv_bin).parent)

            process = await asyncio.create_subprocess_shell(
                command,
                cwd=str(project_root),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=clean_env,
                **get_subprocess_kwargs(),
            )

            async def _read_capped(stream, max_bytes, chunks):
                """Read up to max_bytes from stream into chunks list."""
                total = 0
                while True:
                    chunk = await stream.read(8192)
                    if not chunk:
                        break
                    remaining = max_bytes - total
                    if remaining > 0:
                        chunks.append(chunk[:remaining])
                    total += len(chunk)
                return total > max_bytes

            stdout_trunc, stderr_trunc = await asyncio.wait_for(
                asyncio.gather(
                    _read_capped(process.stdout, MAX_BASH_OUTPUT, stdout_chunks),
                    _read_capped(process.stderr, MAX_BASH_OUTPUT, stderr_chunks),
                ),
                timeout=timeout,
            )

            try:
                await asyncio.wait_for(process.wait(), timeout=5)
            except asyncio.TimeoutError:
                kill_process_tree(process.pid)
                try:
                    await asyncio.wait_for(process.wait(), timeout=5)
                except asyncio.TimeoutError:
                    pass

            return {
                "stdout": b"".join(stdout_chunks).decode(errors="replace"),
                "stderr": b"".join(stderr_chunks).decode(errors="replace"),
                "exit_code": process.returncode,
                "truncated": stdout_trunc or stderr_trunc,
            }
        except asyncio.TimeoutError:
            kill_process_tree(process.pid)
            try:
                await asyncio.wait_for(process.wait(), timeout=5)
            except asyncio.TimeoutError:
                pass  # Process unreapable — OS will clean up

            partial_stdout = b"".join(stdout_chunks).decode(errors="replace")
            partial_stderr = b"".join(stderr_chunks).decode(errors="replace")
            timeout_msg = f"Command timed out after {timeout}s"

            return {
                "stdout": partial_stdout,
                "stderr": (partial_stderr + "\n" + timeout_msg).strip()
                    if partial_stderr else timeout_msg,
                "exit_code": -1,
                "truncated": False,
            }
        finally:
            self._active_bash_count -= 1

    # ========== Session handlers ==========

    # Priority order for auto-selection (lower index = higher priority)
    _ENV_TYPE_PRIORITY = ["venv", "conda", "system"]

    def _auto_select_kernel(
        self, notebook_path: str
    ) -> tuple[Optional[str], Optional[str], Optional[str]]:
        """
        Auto-select the best kernel for a notebook.

        Priority 1: Match notebook's saved kernelspec metadata against available envs.
        Priority 2: Best available env by type (venv > conda > system).

        Returns:
            (kernel_name, python_path, env_name) — any may be None if no match.
        """
        # Get available environments (cached, fast)
        try:
            envs = self.env_manager.get_available_environments()
        except Exception:
            logger.debug("Failed to get environments for auto-selection")
            return (None, None, None)

        # Filter to envs with ipykernel installed
        ready_envs = [e for e in envs if e.get("has_ipykernel")]
        if not ready_envs:
            return (None, None, None)

        # Priority 1: Try matching notebook's saved kernelspec
        matched = None
        try:
            nb = read_notebook(notebook_path)
            spec = nb.get("metadata", {}).get("kernelspec", {})
            spec_display = spec.get("display_name", "")
            spec_name = spec.get("name", "")

            if spec_display or spec_name:
                for env in ready_envs:
                    # Exact match on display_name to avoid false positives
                    # (e.g. "Python" matching "Python 3.10 (myenv)")
                    if spec_display and env.get("display_name", "") == spec_display:
                        matched = env
                        break
                    # Fallback: match by kernelspec name
                    if spec_name and env.get("name") == spec_name:
                        matched = env
                        break
        except (FileNotFoundError, Exception):
            # Notebook doesn't exist yet or invalid — skip metadata matching
            pass

        # Priority 2: Best by type (venv > conda > system)
        if not matched:
            for env_type in self._ENV_TYPE_PRIORITY:
                matched = next(
                    (e for e in ready_envs if e.get("type") == env_type), None
                )
                if matched:
                    break

        if not matched:
            # Use first ready env as last resort
            matched = ready_envs[0]

        return (
            "python3",
            matched.get("python_path"),
            matched.get("display_name"),
        )

    async def _handle_session_create(self, params: dict) -> dict:
        """
        Create or get a session for a notebook, starting a kernel if needed.

        When kernel_name and python_path are both absent, auto-selects the best
        environment: first tries matching the notebook's saved kernelspec metadata,
        then falls back to priority order (venv > conda > system).

        Params:
            notebook_path: Path to the notebook file
            kernel_name: Jupyter kernel spec name (default: python3)
            python_path: Optional path to Python executable
            env_name: Optional display name for the environment

        Returns:
            session_id: The Jupyter session ID
            kernel_id: The kernel ID
            created: True if new session, False if existing
            kernel_name: The kernel spec name used
            python_path: The Python executable path (if any)
            env_name: The environment display name (if any)
        """
        notebook_path = params.get("notebook_path")
        if not notebook_path or not isinstance(notebook_path, str):
            raise ValueError("notebook_path is required and must be a string")

        kernel_name = params.get("kernel_name")
        python_path = params.get("python_path")
        env_name = params.get("env_name")

        # Auto-select best environment when caller doesn't specify
        if not kernel_name and not python_path:
            kernel_name, python_path, env_name = self._auto_select_kernel(
                notebook_path
            )

        # Default kernel_name if still unset
        if not kernel_name:
            # Check if environments exist but none have ipykernel
            try:
                envs = self.env_manager.get_available_environments()
                if envs and not any(e.get("has_ipykernel") for e in envs):
                    raise ValueError(
                        "No environments with ipykernel installed. "
                        "Run 'pip install ipykernel' in your environment."
                    )
            except ValueError:
                raise
            except Exception:
                pass  # env detection failed — fall through to python3 default
            kernel_name = "python3"

        # get_or_create_session handles the check internally with proper locking
        # and returns whether a new session was created
        session, created = await self.session_manager.get_or_create_session(
            notebook_path=notebook_path,
            kernel_name=kernel_name,
            python_path=python_path,
            env_name=env_name,
        )

        # Notify frontend so kernel.output notifications can route by session_id
        # from the very first message (same WebSocket channel guarantees ordering).
        if created:
            await self.send_notification("session.created", {
                "notebook_path": notebook_path,
                "session_id": session.session_id,
                "kernel_id": session.kernel_id,
                "kernel_name": session.kernel_name,
                "python_path": session.python_path,
                "env_name": session.env_name,
            })

        return {
            "session_id": session.session_id,
            "kernel_id": session.kernel_id,
            "created": created,
            "kernel_name": session.kernel_name,
            "python_path": session.python_path,
            "env_name": session.env_name,
        }

    async def _handle_session_get(self, params: dict) -> dict:
        """
        Get session info for a notebook.

        Params:
            notebook_path: Path to the notebook file

        Returns:
            session: Session info dict or null if no session exists
        """
        notebook_path = params.get("notebook_path")
        if not notebook_path or not isinstance(notebook_path, str):
            raise ValueError("notebook_path is required and must be a string")

        session = self.session_manager.get_session_for_path(notebook_path)

        if session is None:
            return {"session": None}

        return {
            "session": {
                "session_id": session.session_id,
                "kernel_id": session.kernel_id,
                "notebook_path": self.file_manager._normalize_output_path(session.notebook_path),
                "created_at": session.created_at.isoformat(),
                "kernel_name": session.kernel_name,
                "python_path": session.python_path,
                "env_name": session.env_name,
            }
        }

    async def _handle_session_list(self, params: dict) -> dict:
        """
        List all active sessions.

        Returns:
            sessions: List of session info dicts
        """
        return {"sessions": self.session_manager.list_sessions()}

    async def _handle_session_close(self, params: dict) -> dict:
        """
        Close a session and its associated kernel.

        Params:
            notebook_path: Path to the notebook file

        Returns:
            success: True if closed, False if no session existed
        """
        notebook_path = params.get("notebook_path")
        if not notebook_path or not isinstance(notebook_path, str):
            raise ValueError("notebook_path is required and must be a string")

        success = await self.session_manager.close_session(notebook_path)
        return {"success": success}

    async def _handle_session_connect_to_kernel(self, params: dict) -> dict:
        """
        Connect a notebook to an existing kernel (kernel sharing).

        This allows multiple notebooks to share the same kernel state.

        Params:
            notebook_path: Path to the notebook file
            kernel_id: ID of existing kernel to connect to

        Returns:
            session_id: The session ID
            kernel_id: The kernel ID
            shared: True (indicates this is a shared connection)
        """
        notebook_path = params.get("notebook_path")
        if not notebook_path or not isinstance(notebook_path, str):
            raise ValueError("notebook_path is required and must be a string")

        kernel_id = params.get("kernel_id")
        if not kernel_id or not isinstance(kernel_id, str):
            raise ValueError("kernel_id is required and must be a string")

        session = await self.session_manager.connect_to_kernel(
            notebook_path=notebook_path,
            kernel_id=kernel_id,
        )

        return {
            "session_id": session.session_id,
            "kernel_id": session.kernel_id,
            "shared": True,
        }

    async def _handle_session_update_path(self, params: dict) -> dict:
        """
        Update session path after a file/folder rename.

        Params:
            old_path: Previous notebook or folder path
            new_path: New notebook or folder path
            is_folder: Whether this is a folder rename (default: False)

        Returns:
            success: True
            updated: Whether any session was actually renamed
        """
        old_path = params.get("old_path")
        new_path = params.get("new_path")
        if not old_path or not isinstance(old_path, str):
            raise ValueError("old_path is required and must be a string")
        if not new_path or not isinstance(new_path, str):
            raise ValueError("new_path is required and must be a string")

        is_folder = params.get("is_folder", False)
        updated = self.session_manager.rename_session(
            old_path, new_path, is_folder=is_folder
        )
        return {"success": True, "updated": updated}
