venus_server/

session.rs

//! Notebook session management.
//!
//! Manages the state of an active notebook session including
//! compilation, execution, and output caching.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use tokio::sync::{RwLock, broadcast};
use venus::widgets::{WidgetDef, WidgetValue};
use venus_core::compile::{
    CellCompiler, CompilationResult, CompilerConfig, ToolchainManager, UniverseBuilder,
};
use venus_core::execute::{ExecutorKillHandle, ProcessExecutor};
use venus_core::graph::{CellId, CellInfo, CellParser, GraphEngine, MarkdownCell, MoveDirection, SourceEditor};
use venus_core::paths::NotebookDirs;

use crate::error::{ServerError, ServerResult};
use crate::protocol::{CellOutput, CellState, CellStatus, DependencyEdge, ServerMessage};
use crate::undo::{UndoManager, UndoableOperation};
use venus_core::state::BoxedOutput;

/// Shared interrupt flag that can be checked without locks.
pub type InterruptFlag = Arc<AtomicBool>;

/// Capacity for the broadcast channel.
/// 256 messages should be sufficient for normal notebook operation.
/// If clients fall behind, older messages will be dropped.
const MESSAGE_CHANNEL_CAPACITY: usize = 256;

/// A notebook session.
pub struct NotebookSession {
    /// Path to the notebook file.
    path: PathBuf,

    /// Parsed code cells.
    cells: Vec<CellInfo>,

    /// Parsed markdown cells.
    markdown_cells: Vec<MarkdownCell>,

    /// Dependency graph.
    graph: GraphEngine,

    /// Cell states for clients (both code and markdown).
    cell_states: HashMap<CellId, CellState>,

    /// Toolchain manager.
    toolchain: ToolchainManager,

    /// Compiler configuration.
    config: CompilerConfig,

    /// Universe path (compiled dependencies).
    universe_path: Option<PathBuf>,

    /// Dependencies hash for cache invalidation.
    deps_hash: u64,

    /// Broadcast channel for server messages.
    tx: broadcast::Sender<ServerMessage>,

    /// Whether an execution is in progress.
    executing: bool,

    /// Cached cell outputs for dependency passing.
    /// Maps cell ID to its serialized output.
    cell_outputs: HashMap<CellId, Arc<BoxedOutput>>,

    /// Process-based executor for isolated cell execution.
    /// Uses worker processes that can be killed for true interruption.
    executor: ProcessExecutor,

    /// Optional execution timeout for execute_all.
    /// After this duration, the executor kills the current worker.
    execution_timeout: Option<Duration>,

    /// Shared flag indicating if current execution was interrupted by user.
    /// When true, errors should be reported as "interrupted" not as failures.
    /// This is shared with AppState so interrupt handler can set it.
    interrupted: InterruptFlag,

    /// Widget values per cell.
    /// Maps cell ID -> widget ID -> current value.
    widget_values: HashMap<CellId, HashMap<String, WidgetValue>>,

    /// Widget definitions per cell (from last execution).
    /// Used to send widget state to newly connected clients.
    widget_defs: HashMap<CellId, Vec<WidgetDef>>,

    /// Execution history per cell.
    /// Stores both serialized output (for dependent cells) and display output.
    cell_output_history: HashMap<CellId, Vec<OutputHistoryEntry>>,

    /// Current history index per cell.
    cell_history_index: HashMap<CellId, usize>,

    /// Undo/redo manager for cell operations.
    undo_manager: UndoManager,
}

/// Maximum number of history entries per cell.
const MAX_HISTORY_PER_CELL: usize = 10;

/// A single history entry for a cell's execution.
#[derive(Clone)]
pub struct OutputHistoryEntry {
    /// Serialized output for passing to dependent cells.
    pub serialized: Arc<BoxedOutput>,
    /// Display output for the frontend.
    pub display: CellOutput,
    /// Timestamp when this execution completed.
    pub timestamp: u64,
}

/// Thread-safe session handle.
pub type SessionHandle = Arc<RwLock<NotebookSession>>;

impl NotebookSession {
    /// Create a new notebook session.
    ///
    /// Uses process isolation for cell execution, allowing true interruption
    /// by killing worker processes.
    ///
    /// The `interrupted` flag is shared with AppState so the interrupt handler
    /// can set it without needing the session lock.
    pub fn new(
        path: impl AsRef<Path>,
        interrupted: InterruptFlag,
    ) -> ServerResult<(Self, broadcast::Receiver<ServerMessage>)> {
        let path = path.as_ref().canonicalize().map_err(|e| ServerError::Io {
            path: path.as_ref().to_path_buf(),
            message: e.to_string(),
        })?;

        // Set up directories using shared abstraction
        let dirs = NotebookDirs::from_notebook_path(&path)?;

        let toolchain = ToolchainManager::new()?;
        let config = CompilerConfig::for_notebook(&dirs);

        let (tx, rx) = broadcast::channel(MESSAGE_CHANNEL_CAPACITY);

        // Create process executor with warm worker pool
        let executor = ProcessExecutor::new(&dirs.state_dir)?;

        let mut session = Self {
            path,
            cells: Vec::new(),
            markdown_cells: Vec::new(),
            graph: GraphEngine::new(),
            cell_states: HashMap::new(),
            toolchain,
            config,
            universe_path: None,
            deps_hash: 0,
            tx,
            executing: false,
            cell_outputs: HashMap::new(),
            executor,
            execution_timeout: None,
            interrupted,
            widget_values: HashMap::new(),
            widget_defs: HashMap::new(),
            cell_output_history: HashMap::new(),
            cell_history_index: HashMap::new(),
            undo_manager: UndoManager::new(),
        };

        session.reload()?;

        Ok((session, rx))
    }

    /// Get the notebook path.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Subscribe to server messages.
    pub fn subscribe(&self) -> broadcast::Receiver<ServerMessage> {
        self.tx.subscribe()
    }

    /// Get a cell by ID.
    fn get_cell(&self, cell_id: CellId) -> Option<&CellInfo> {
        self.cells.iter().find(|c| c.id == cell_id)
    }

    /// Set the status of a code cell.
    fn set_cell_status(&mut self, cell_id: CellId, status: CellStatus) {
        if let Some(CellState::Code { status: cell_status, .. }) = self.cell_states.get_mut(&cell_id) {
            *cell_status = status;
        }
    }

    /// Broadcast a server message, ignoring send failures.
    pub fn broadcast(&self, msg: ServerMessage) {
        let _ = self.tx.send(msg);
    }

    /// Reload the notebook from disk.
    pub fn reload(&mut self) -> ServerResult<()> {
        let source = std::fs::read_to_string(&self.path)?;

        // Parse cells (both code and markdown)
        let mut parser = CellParser::new();
        let parse_result = parser.parse_file(&self.path)?;
        self.cells = parse_result.code_cells;
        self.markdown_cells = parse_result.markdown_cells;

        // Build graph and update code cells with real IDs (parser returns placeholder IDs)
        self.graph = GraphEngine::new();
        for cell in &mut self.cells {
            let real_id = self.graph.add_cell(cell.clone());
            cell.id = real_id;
        }
        self.graph.resolve_dependencies()?;

        // Assign unique IDs to markdown cells (they don't participate in the dependency graph)
        let mut next_id = if let Some(max_code_id) = self.cells.iter().map(|c| c.id.as_usize()).max() {
            max_code_id + 1
        } else {
            0
        };
        for md_cell in &mut self.markdown_cells {
            md_cell.id = CellId::new(next_id);
            next_id += 1;
        }

        // Build universe (always needed for bincode/serde runtime)
        let mut universe_builder =
            UniverseBuilder::new(self.config.clone(), self.toolchain.clone());
        universe_builder.parse_dependencies(&source)?;

        self.universe_path = Some(universe_builder.build()?);
        self.deps_hash = universe_builder.deps_hash();

        // Update cell states
        self.update_cell_states();

        // Broadcast full state update (includes all cells, not just graph edges)
        let state = self.get_state();
        self.broadcast(state);

        Ok(())
    }

    /// Strip the first heading from a doc comment (since it's used as display name).
    ///
    /// If the doc comment starts with `# Heading`, removes that line and returns
    /// the rest of the content. Otherwise returns the original content.
    fn strip_display_name_from_description(doc_comment: &Option<String>) -> Option<String> {
        doc_comment.as_ref().map(|doc| {
            let lines: Vec<&str> = doc.lines().collect();

            // Find the first heading line
            if let Some(first_line) = lines.first() {
                let trimmed = first_line.trim();
                if trimmed.starts_with('#') {
                    // Skip the heading line and return the rest
                    let remaining: Vec<&str> = lines.iter().skip(1)
                        .map(|s| *s)
                        .collect();

                    // Trim leading empty lines
                    let trimmed_lines: Vec<&str> = remaining.iter()
                        .skip_while(|line| line.trim().is_empty())
                        .map(|s| *s)
                        .collect();

                    if trimmed_lines.is_empty() {
                        return None;
                    }

                    return Some(trimmed_lines.join("\n"));
                }
            }

            // No heading found, return original
            Some(doc.clone())
        }).flatten()
    }

    /// Update cell states from parsed cells.
    fn update_cell_states(&mut self) {
        let mut new_states = HashMap::new();

        // Add code cells
        for cell in &self.cells {
            let existing = self.cell_states.get(&cell.id);

            // Extract status, output, dirty from existing state if it's a code cell
            let (status, output, dirty) = if let Some(CellState::Code { status, output, dirty, .. }) = existing {
                (*status, output.clone(), *dirty)
            } else {
                (CellStatus::default(), None, true)
            };

            let state = CellState::Code {
                id: cell.id,
                name: cell.name.clone(),
                display_name: cell.display_name.clone(),
                source: cell.source_code.clone(),
                description: Self::strip_display_name_from_description(&cell.doc_comment),
                return_type: cell.return_type.clone(),
                dependencies: cell
                    .dependencies
                    .iter()
                    .map(|d| d.param_name.clone())
                    .collect(),
                status,
                output,
                dirty,
            };
            new_states.insert(cell.id, state);
        }

        // Add markdown cells
        for md_cell in &self.markdown_cells {
            let state = CellState::Markdown {
                id: md_cell.id,
                content: md_cell.content.clone(),
            };
            new_states.insert(md_cell.id, state);
        }

        self.cell_states = new_states;
    }

    /// Broadcast current graph state.
    fn broadcast_graph_update(&self) {
        let edges: Vec<DependencyEdge> = self
            .cells
            .iter()
            .flat_map(|cell| {
                cell.dependencies.iter().filter_map(|dep| {
                    self.cells
                        .iter()
                        .find(|c| c.name == dep.param_name)
                        .map(|producer| DependencyEdge {
                            from: producer.id,
                            to: cell.id,
                            param_name: dep.param_name.clone(),
                        })
                })
            })
            .collect();

        let order = match self.graph.topological_order() {
            Ok(order) => order,
            Err(e) => {
                tracing::error!("Failed to compute topological order: {}", e);
                Vec::new()
            }
        };
        let levels = self.graph.topological_levels(&order);

        self.broadcast(ServerMessage::GraphUpdated { edges, levels });
    }

    /// Get the full notebook state.
    pub fn get_state(&self) -> ServerMessage {
        // Source order: all cells (code + markdown) in the order they appear in the .rs file
        let mut all_cells: Vec<(CellId, usize)> = Vec::new();

        // Add code cells with their line numbers
        for cell in &self.cells {
            all_cells.push((cell.id, cell.span.start_line));
        }

        // Add markdown cells with their line numbers
        for md_cell in &self.markdown_cells {
            all_cells.push((md_cell.id, md_cell.span.start_line));
        }

        // Sort by line number
        all_cells.sort_by_key(|(_, line)| *line);

        let source_order: Vec<CellId> = all_cells.into_iter().map(|(id, _)| id).collect();

        // Execution order: topologically sorted for dependency resolution (code cells only)
        let execution_order = match self.graph.topological_order() {
            Ok(order) => order,
            Err(e) => {
                tracing::error!("Failed to compute execution order: {}", e);
                Vec::new()
            }
        };

        ServerMessage::NotebookState {
            path: self.path.display().to_string(),
            cells: self.cell_states.values().cloned().collect(),
            source_order,
            execution_order,
        }
    }

    /// Execute a specific cell.
    ///
    /// Uses process isolation - the cell runs in a worker process that can
    /// be killed immediately for interruption.
    pub async fn execute_cell(&mut self, cell_id: CellId) -> ServerResult<()> {
        if self.executing {
            return Err(ServerError::ExecutionInProgress);
        }

        let cell = self
            .get_cell(cell_id)
            .ok_or(ServerError::CellNotFound(cell_id))?
            .clone();

        self.executing = true;

        // Check if all dependencies have outputs available
        let missing_deps: Vec<&str> = cell
            .dependencies
            .iter()
            .filter(|dep| {
                let producer = self.cells.iter().find(|c| c.name == dep.param_name);
                match producer {
                    Some(c) => !self.cell_outputs.contains_key(&c.id),
                    None => true,
                }
            })
            .map(|d| d.param_name.as_str())
            .collect();

        if !missing_deps.is_empty() {
            self.set_cell_status(cell_id, CellStatus::Error);
            self.broadcast(ServerMessage::CellError {
                cell_id,
                error: format!(
                    "Missing dependencies: {}. Run dependent cells first.",
                    missing_deps.join(", ")
                ),
                location: None,
            });
            self.executing = false;
            return Ok(());
        }

        // Compile
        self.set_cell_status(cell_id, CellStatus::Compiling);

        let mut compiler = CellCompiler::new(self.config.clone(), self.toolchain.clone());
        if let Some(ref up) = self.universe_path {
            compiler = compiler.with_universe(up.clone());
        }

        let result = compiler.compile(&cell, self.deps_hash);

        match result {
            CompilationResult::Success(compiled) | CompilationResult::Cached(compiled) => {
                // Execute
                self.set_cell_status(cell_id, CellStatus::Running);
                self.broadcast(ServerMessage::CellStarted { cell_id });

                let start = Instant::now();

                // Register the compiled cell with the executor
                self.executor.register_cell(compiled, cell.dependencies.len());

                // Gather dependency outputs in the order the cell expects them
                let inputs: Vec<Arc<BoxedOutput>> = cell
                    .dependencies
                    .iter()
                    .filter_map(|dep| {
                        self.cells
                            .iter()
                            .find(|c| c.name == dep.param_name)
                            .and_then(|c| self.cell_outputs.get(&c.id).cloned())
                    })
                    .collect();

                // Get ALL widget values from all cells (widgets can be in any cell)
                let widget_values = self.get_all_widget_values();
                let widget_values_json = if widget_values.is_empty() {
                    Vec::new()
                } else {
                    serde_json::to_vec(&widget_values).unwrap_or_default()
                };

                // Execute the cell in an isolated worker process with widget values
                let exec_result = self.executor.execute_cell_with_widgets(
                    cell_id,
                    &inputs,
                    widget_values_json,
                );

                let duration = start.elapsed();

                match exec_result {
                    Ok((output, widgets_json)) => {
                        // Store output for dependent cells
                        let output_arc = Arc::new(output);
                        self.cell_outputs.insert(cell_id, output_arc.clone());

                        // Also store in executor state for consistency
                        self.executor.state_mut().store_output(cell_id, (*output_arc).clone());

                        // Parse and store widget definitions
                        let widgets: Vec<WidgetDef> = if widgets_json.is_empty() {
                            Vec::new()
                        } else {
                            serde_json::from_slice(&widgets_json).unwrap_or_default()
                        };
                        self.store_widget_defs(cell_id, widgets.clone());

                        let cell_output = CellOutput {
                            text: output_arc.display_text().map(|s| s.to_string()),
                            html: None,
                            image: None,
                            json: None,
                            widgets,
                        };

                        // Add to history
                        self.add_to_history(cell_id, output_arc.clone(), cell_output.clone());

                        if let Some(state) = self.cell_states.get_mut(&cell_id) {
                            state.set_status(CellStatus::Success);
                            state.set_output(Some(cell_output.clone()));
                            state.set_dirty(false);
                        }

                        self.broadcast(ServerMessage::CellCompleted {
                            cell_id,
                            duration_ms: duration.as_millis() as u64,
                            output: Some(cell_output),
                        });
                    }
                    Err(e) => {
                        // Check if this was an abort or user-initiated interrupt
                        let was_interrupted = self.interrupted.swap(false, Ordering::SeqCst);
                        if matches!(e, venus_core::Error::Aborted) || was_interrupted {
                            // Send friendly "interrupted" message instead of error
                            self.set_cell_status(cell_id, CellStatus::Idle);
                            self.broadcast(ServerMessage::ExecutionAborted { cell_id: Some(cell_id) });
                        } else {
                            self.set_cell_status(cell_id, CellStatus::Error);
                            self.broadcast(ServerMessage::CellError {
                                cell_id,
                                error: e.to_string(),
                                location: None,
                            });
                        }
                    }
                }
            }
            CompilationResult::Failed { errors, .. } => {
                self.set_cell_status(cell_id, CellStatus::Error);

                let compile_errors = errors
                    .iter()
                    .map(|e| crate::protocol::CompileErrorInfo {
                        message: e.message.clone(),
                        code: e.code.clone(),
                        location: e.spans.first().map(|s| crate::protocol::SourceLocation {
                            line: s.location.line as u32,
                            column: s.location.column as u32,
                            end_line: s.end_location.as_ref().map(|l| l.line as u32),
                            end_column: s.end_location.as_ref().map(|l| l.column as u32),
                        }),
                        rendered: e.rendered.clone(),
                    })
                    .collect();

                self.broadcast(ServerMessage::CompileError {
                    cell_id,
                    errors: compile_errors,
                });
            }
        }

        self.executing = false;
        Ok(())
    }

    /// Execute all cells in order.
    ///
    /// If `execution_timeout` is set, kills the worker process after that duration.
    /// Unlike cooperative cancellation, this immediately terminates the cell.
    pub async fn execute_all(&mut self) -> ServerResult<()> {
        let order = self.graph.topological_order()?;
        let start = Instant::now();
        let timeout = self.execution_timeout;

        for cell_id in order {
            // Check timeout before each cell
            if timeout.is_some_and(|max_duration| start.elapsed() > max_duration) {
                self.executor.abort();
                self.broadcast(ServerMessage::ExecutionAborted { cell_id: Some(cell_id) });
                return Err(ServerError::ExecutionTimeout);
            }

            self.execute_cell(cell_id).await?;
        }
        Ok(())
    }

    /// Mark a cell as dirty (needs re-execution).
    pub fn mark_dirty(&mut self, cell_id: CellId) {
        if let Some(state) = self.cell_states.get_mut(&cell_id) {
            state.set_dirty(true);
        }

        // Also mark dependents as dirty
        let dependents = self.graph.invalidated_cells(cell_id);
        for dep_id in dependents {
            if let Some(state) = self.cell_states.get_mut(&dep_id) {
                state.set_dirty(true);
            }
        }
    }

    /// Check if execution is in progress.
    pub fn is_executing(&self) -> bool {
        self.executing
    }

    /// Abort the current execution immediately.
    ///
    /// Unlike cooperative cancellation, this **kills the worker process**,
    /// providing true interruption even for long-running computations.
    /// Returns `true` if there was an execution in progress to abort.
    pub fn abort(&mut self) -> bool {
        if self.executing {
            // Kill the worker process - this is immediate
            self.executor.abort();
            self.broadcast(ServerMessage::ExecutionAborted { cell_id: None });
            self.executing = false;
            true
        } else {
            false
        }
    }

    /// Set the execution timeout for execute_all.
    ///
    /// When set, execute_all will kill the worker process after this duration,
    /// providing immediate interruption of even long-running cells.
    pub fn set_execution_timeout(&mut self, timeout: Option<Duration>) {
        self.execution_timeout = timeout;
    }

    /// Get the current execution timeout.
    pub fn execution_timeout(&self) -> Option<Duration> {
        self.execution_timeout
    }

    /// Set the interrupted flag.
    ///
    /// When true, execution errors will be reported as "interrupted"
    /// rather than as failures, showing a friendly message to users.
    pub fn set_interrupted(&mut self, value: bool) {
        self.interrupted.store(value, Ordering::SeqCst);
    }

    /// Get a kill handle for the executor.
    ///
    /// This handle can be used from another task to kill the current execution
    /// without needing to acquire the session lock.
    pub fn get_kill_handle(&self) -> Option<ExecutorKillHandle> {
        self.executor.get_kill_handle()
    }

    /// Restart the kernel: kill WorkerPool, spin up new one, clear memory state, preserve source.
    ///
    /// This clears all execution state including:
    /// - Cell outputs and output history
    /// - Widget values
    /// - Cached serialized outputs
    /// - Cell execution status (all cells reset to Idle)
    ///
    /// Source code and cell definitions are preserved.
    pub fn restart_kernel(&mut self) -> ServerResult<()> {
        // Abort any running execution first
        if self.executing {
            self.abort();
        }

        // Reload notebook from disk (picks up any file changes)
        self.reload()?;

        // Shutdown old executor and worker pool
        self.executor.shutdown();

        // Reconstruct state directory path
        let dirs = NotebookDirs::from_notebook_path(&self.path)?;

        // Create new ProcessExecutor with warm worker pool
        self.executor = ProcessExecutor::new(&dirs.state_dir)?;

        // Clear all execution state
        self.cell_outputs.clear();
        self.widget_values.clear();
        self.widget_defs.clear();
        self.cell_output_history.clear();
        self.cell_history_index.clear();

        // Reset all cell states to Idle and clear outputs
        for state in self.cell_states.values_mut() {
            state.set_status(CellStatus::Idle);
            state.clear_output();
            state.set_dirty(false);
        }

        // Broadcast kernel restarted message
        self.broadcast(ServerMessage::KernelRestarted { error: None });

        // Send updated state to all clients
        let state_msg = self.get_state();
        self.broadcast(state_msg);

        Ok(())
    }

    /// Clear all cell outputs without restarting the kernel.
    ///
    /// This clears the display outputs but preserves:
    /// - Worker pool and execution state
    /// - Widget values
    /// - Cell source code
    ///
    /// All cells are marked as needing re-execution (dirty).
    pub fn clear_outputs(&mut self) {
        // Clear outputs from cell states
        for state in self.cell_states.values_mut() {
            state.clear_output();
            state.set_dirty(true); // Mark as needing re-execution
        }

        // Clear output history
        self.cell_output_history.clear();
        self.cell_history_index.clear();

        // Broadcast outputs cleared message
        self.broadcast(ServerMessage::OutputsCleared { error: None });

        // Send updated state to all clients
        let state_msg = self.get_state();
        self.broadcast(state_msg);
    }

    /// Get IDs of all dirty cells in topological order.
    pub fn get_dirty_cell_ids(&self) -> Vec<CellId> {
        let order = match self.graph.topological_order() {
            Ok(order) => order,
            Err(e) => {
                tracing::error!("Failed to compute order for dirty cells: {}", e);
                Vec::new()
            }
        };
        order
            .into_iter()
            .filter(|id| self.cell_states.get(id).is_some_and(|state| state.is_dirty()))
            .collect()
    }

    /// Update a widget value for a cell.
    ///
    /// This stores the new value but does NOT trigger re-execution.
    /// The user must explicitly run the cell to see the effect.
    pub fn update_widget_value(&mut self, cell_id: CellId, widget_id: String, value: WidgetValue) {
        self.widget_values
            .entry(cell_id)
            .or_default()
            .insert(widget_id, value);
    }

    /// Get widget values for a cell.
    pub fn get_widget_values(&self, cell_id: CellId) -> HashMap<String, WidgetValue> {
        self.widget_values
            .get(&cell_id)
            .cloned()
            .unwrap_or_default()
    }

    /// Get ALL widget values from all cells, flattened into a single map.
    /// Widget IDs should be unique across the notebook.
    pub fn get_all_widget_values(&self) -> HashMap<String, WidgetValue> {
        let mut all_values = HashMap::new();
        for cell_widgets in self.widget_values.values() {
            for (widget_id, value) in cell_widgets {
                all_values.insert(widget_id.clone(), value.clone());
            }
        }
        all_values
    }

    /// Get widget definitions for a cell.
    pub fn get_widget_defs(&self, cell_id: CellId) -> Vec<WidgetDef> {
        self.widget_defs
            .get(&cell_id)
            .cloned()
            .unwrap_or_default()
    }

    /// Store widget definitions from cell execution.
    fn store_widget_defs(&mut self, cell_id: CellId, widgets: Vec<WidgetDef>) {
        if widgets.is_empty() {
            self.widget_defs.remove(&cell_id);
        } else {
            self.widget_defs.insert(cell_id, widgets);
        }
    }

    /// Add an execution result to history.
    fn add_to_history(&mut self, cell_id: CellId, serialized: Arc<BoxedOutput>, display: CellOutput) {
        use std::time::{SystemTime, UNIX_EPOCH};

        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;

        let entry = OutputHistoryEntry {
            serialized,
            display,
            timestamp,
        };

        let history = self.cell_output_history.entry(cell_id).or_default();
        history.push(entry);

        // Trim if too long
        while history.len() > MAX_HISTORY_PER_CELL {
            history.remove(0);
        }

        // Set current index to the latest entry
        self.cell_history_index.insert(cell_id, history.len() - 1);
    }

    /// Select a history entry for a cell, making it the current output.
    /// Returns the display output if successful.
    pub fn select_history_entry(&mut self, cell_id: CellId, index: usize) -> Option<CellOutput> {
        // Clone what we need before doing any mutations (to avoid borrow conflicts)
        let (serialized, display) = {
            let history = self.cell_output_history.get(&cell_id)?;
            let entry = history.get(index)?;
            (entry.serialized.clone(), entry.display.clone())
        };

        // Update the current output for dependent cells
        self.cell_outputs.insert(cell_id, serialized.clone());
        self.executor.state_mut().store_output(cell_id, (*serialized).clone());

        // Update the cell state
        if let Some(state) = self.cell_states.get_mut(&cell_id) {
            state.set_output(Some(display.clone()));
        }

        // Update history index
        self.cell_history_index.insert(cell_id, index);

        // Mark dependent cells as dirty
        self.mark_dependents_dirty(cell_id);

        Some(display)
    }

    /// Mark all cells that depend on the given cell as dirty.
    fn mark_dependents_dirty(&mut self, cell_id: CellId) {
        // Use the graph's invalidated_cells which returns all dependents
        let dependents = self.graph.invalidated_cells(cell_id);

        // Skip the first one (the changed cell itself) and mark the rest as dirty
        for dep_id in dependents.into_iter().skip(1) {
            if let Some(state) = self.cell_states.get_mut(&dep_id) {
                state.set_dirty(true);
            }
        }
    }

    /// Get history count for a cell.
    pub fn get_history_count(&self, cell_id: CellId) -> usize {
        self.cell_output_history.get(&cell_id).map(|h| h.len()).unwrap_or(0)
    }

    /// Get current history index for a cell.
    pub fn get_history_index(&self, cell_id: CellId) -> usize {
        self.cell_history_index.get(&cell_id).copied().unwrap_or(0)
    }

    /// Get reference to cell states.
    pub fn cell_states(&self) -> &HashMap<CellId, CellState> {
        &self.cell_states
    }

    /// Insert a new cell after the specified cell.
    ///
    /// Modifies the source file and triggers a reload.
    /// Returns the name of the newly created cell.
    pub fn insert_cell(&mut self, after_cell_id: Option<CellId>) -> ServerResult<String> {
        // Convert CellId to cell name if provided
        let after_name = after_cell_id.and_then(|id| {
            self.cells.iter().find(|c| c.id == id).map(|c| c.name.clone())
        });

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        let new_name = editor.insert_cell(after_name.as_deref())?;
        editor.save()?;

        // Record for undo (with position for redo)
        self.undo_manager.record(UndoableOperation::InsertCell {
            cell_name: new_name.clone(),
            after_cell_name: after_name,
        });

        // File watcher will trigger reload, but we can also reload now
        // to ensure immediate consistency
        self.reload()?;

        Ok(new_name)
    }

    /// Delete a cell from the notebook.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn delete_cell(&mut self, cell_id: CellId) -> ServerResult<()> {
        // Find the cell name
        let cell_name = self.cells
            .iter()
            .find(|c| c.id == cell_id)
            .map(|c| c.name.clone())
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        // Check if any other cells depend on this cell
        let dependents: Vec<String> = self.cells
            .iter()
            .filter(|c| c.id != cell_id) // Don't check self
            .filter(|c| {
                c.dependencies
                    .iter()
                    .any(|dep| dep.param_name == cell_name)
            })
            .map(|c| c.name.clone())
            .collect();

        if !dependents.is_empty() {
            return Err(ServerError::InvalidOperation(format!(
                "Cannot delete cell '{}' because it is used by: {}",
                cell_name,
                dependents.join(", ")
            )));
        }

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;

        // Capture source and position before deletion (for undo)
        let source = editor.get_cell_source(&cell_name)?;
        let after_cell_name = editor.get_previous_cell_name(&cell_name)?;

        editor.delete_cell(&cell_name)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::DeleteCell {
            cell_name: cell_name.clone(),
            source,
            after_cell_name,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Duplicate a cell in the notebook.
    ///
    /// Creates a copy of the cell with a unique name.
    /// Returns the name of the new cell.
    pub fn duplicate_cell(&mut self, cell_id: CellId) -> ServerResult<String> {
        // Find the cell name
        let cell_name = self.cells
            .iter()
            .find(|c| c.id == cell_id)
            .map(|c| c.name.clone())
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        let new_name = editor.duplicate_cell(&cell_name)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::DuplicateCell {
            original_cell_name: cell_name,
            new_cell_name: new_name.clone(),
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(new_name)
    }

    /// Move a cell up or down in the notebook.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn move_cell(&mut self, cell_id: CellId, direction: MoveDirection) -> ServerResult<()> {
        // Find the cell name
        let cell_name = self.cells
            .iter()
            .find(|c| c.id == cell_id)
            .map(|c| c.name.clone())
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.move_cell(&cell_name, direction)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::MoveCell {
            cell_name,
            direction,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Rename a cell's display name.
    ///
    /// Updates the cell's doc comment with the new display name and reloads the notebook.
    pub fn rename_cell(&mut self, cell_id: CellId, new_display_name: String) -> ServerResult<()> {
        // Find the cell name and current display name
        let (cell_name, old_display_name) = self.cells
            .iter()
            .find(|c| c.id == cell_id)
            .map(|c| (c.name.clone(), c.display_name.clone()))
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.rename_cell(&cell_name, &new_display_name)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::RenameCell {
            cell_name,
            old_display_name,
            new_display_name,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Insert a new markdown cell.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn insert_markdown_cell(&mut self, content: String, after_cell_id: Option<CellId>) -> ServerResult<()> {
        // Convert cell ID to line number if provided
        let after_line = after_cell_id.and_then(|id| {
            // Try to find in code cells
            self.cells.iter().find(|c| c.id == id)
                .map(|c| c.span.end_line)
                .or_else(|| {
                    // Try to find in markdown cells
                    self.markdown_cells.iter().find(|m| m.id == id)
                        .map(|m| m.span.end_line)
                })
        });

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.insert_markdown_cell(&content, after_line)?;

        // Get the line range of the newly inserted cell (approximate)
        let start_line = after_line.map(|l| l + 1).unwrap_or(0);
        let line_count = content.lines().count();
        let end_line = start_line + line_count;

        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::InsertMarkdownCell {
            start_line,
            end_line,
            content: content.clone(),
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Edit a markdown cell's content.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn edit_markdown_cell(&mut self, cell_id: CellId, new_content: String) -> ServerResult<()> {
        // Find the markdown cell
        let md_cell = self.markdown_cells
            .iter()
            .find(|m| m.id == cell_id)
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        let start_line = md_cell.span.start_line;
        let end_line = md_cell.span.end_line;
        let old_content = md_cell.content.clone();
        let is_module_doc = md_cell.is_module_doc;

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.edit_markdown_cell(start_line, end_line, &new_content, is_module_doc)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::EditMarkdownCell {
            start_line,
            end_line,
            old_content,
            new_content,
            is_module_doc,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Delete a markdown cell.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn delete_markdown_cell(&mut self, cell_id: CellId) -> ServerResult<()> {
        // Find the markdown cell
        let md_cell = self.markdown_cells
            .iter()
            .find(|m| m.id == cell_id)
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        let start_line = md_cell.span.start_line;
        let end_line = md_cell.span.end_line;
        let content = md_cell.content.clone();

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.delete_markdown_cell(start_line, end_line)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::DeleteMarkdownCell {
            start_line,
            content,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Move a markdown cell up or down.
    ///
    /// Modifies the .rs source file and reloads the notebook.
    pub fn move_markdown_cell(&mut self, cell_id: CellId, direction: MoveDirection) -> ServerResult<()> {
        // Find the markdown cell
        let md_cell = self.markdown_cells
            .iter()
            .find(|m| m.id == cell_id)
            .ok_or_else(|| ServerError::CellNotFound(cell_id))?;

        let start_line = md_cell.span.start_line;
        let end_line = md_cell.span.end_line;

        // Load and edit the source file
        let mut editor = SourceEditor::load(&self.path)?;
        editor.move_markdown_cell(start_line, end_line, direction)?;
        editor.save()?;

        // Record for undo
        self.undo_manager.record(UndoableOperation::MoveMarkdownCell {
            start_line,
            end_line,
            direction,
        });

        // Reload to update in-memory state
        self.reload()?;

        Ok(())
    }

    /// Undo the last cell management operation.
    ///
    /// Returns a description of what was undone, or an error if undo failed.
    pub fn undo(&mut self) -> ServerResult<String> {
        let operation = self.undo_manager.pop_undo()
            .ok_or_else(|| ServerError::InvalidOperation("Nothing to undo".to_string()))?;

        let description = operation.undo_description();

        // Execute the reverse operation
        let mut editor = SourceEditor::load(&self.path)?;

        match &operation {
            UndoableOperation::InsertCell { cell_name, .. } => {
                // Undo insert = delete
                editor.delete_cell(cell_name)?;
            }
            UndoableOperation::DeleteCell { source, after_cell_name, .. } => {
                // Undo delete = restore
                editor.restore_cell(source, after_cell_name.as_deref())?;
            }
            UndoableOperation::DuplicateCell { new_cell_name, .. } => {
                // Undo duplicate = delete the new cell
                editor.delete_cell(new_cell_name)?;
            }
            UndoableOperation::MoveCell { cell_name, direction } => {
                // Undo move = move in opposite direction
                let reverse_direction = match direction {
                    MoveDirection::Up => MoveDirection::Down,
                    MoveDirection::Down => MoveDirection::Up,
                };
                editor.move_cell(cell_name, reverse_direction)?;
            }
            UndoableOperation::RenameCell { cell_name, old_display_name, .. } => {
                // Undo rename = restore old display name
                editor.rename_cell(cell_name, old_display_name)?;
            }
            UndoableOperation::InsertMarkdownCell { start_line, end_line, .. } => {
                // Undo insert markdown = delete it
                editor.delete_markdown_cell(*start_line, *end_line)?;
            }
            UndoableOperation::EditMarkdownCell { start_line, end_line, old_content, is_module_doc, .. } => {
                // Undo edit markdown = restore old content
                editor.edit_markdown_cell(*start_line, *end_line, old_content, *is_module_doc)?;
            }
            UndoableOperation::DeleteMarkdownCell { start_line, content } => {
                // Undo delete markdown = restore it
                let after_line = if *start_line > 0 { Some(start_line - 1) } else { None };
                editor.insert_markdown_cell(content, after_line)?;
            }
            UndoableOperation::MoveMarkdownCell { start_line, end_line, direction } => {
                // Undo move markdown = move in opposite direction
                let reverse_direction = match direction {
                    MoveDirection::Up => MoveDirection::Down,
                    MoveDirection::Down => MoveDirection::Up,
                };
                editor.move_markdown_cell(*start_line, *end_line, reverse_direction)?;
            }
        }

        editor.save()?;

        // Record for redo
        self.undo_manager.record_redo(operation);

        // Reload to update in-memory state
        self.reload()?;

        Ok(description)
    }

    /// Redo the last undone operation.
    ///
    /// Returns a description of what was redone, or an error if redo failed.
    pub fn redo(&mut self) -> ServerResult<String> {
        let operation = self.undo_manager.pop_redo()
            .ok_or_else(|| ServerError::InvalidOperation("Nothing to redo".to_string()))?;

        let description = operation.description();

        // Execute the original operation
        let mut editor = SourceEditor::load(&self.path)?;

        match &operation {
            UndoableOperation::InsertCell { after_cell_name, .. } => {
                // Re-insert at the original position
                let _ = editor.insert_cell(after_cell_name.as_deref())?;
            }
            UndoableOperation::DeleteCell { cell_name, .. } => {
                // Redo delete = delete again
                editor.delete_cell(cell_name)?;
            }
            UndoableOperation::DuplicateCell { original_cell_name, .. } => {
                // Redo duplicate = duplicate again (new name will be generated)
                let _ = editor.duplicate_cell(original_cell_name)?;
            }
            UndoableOperation::MoveCell { cell_name, direction } => {
                // Redo move = move in same direction
                editor.move_cell(cell_name, *direction)?;
            }
            UndoableOperation::RenameCell { cell_name, new_display_name, .. } => {
                // Redo rename = apply new display name again
                editor.rename_cell(cell_name, new_display_name)?;
            }
            UndoableOperation::InsertMarkdownCell { start_line, content, .. } => {
                // Redo insert markdown = insert again at original position
                let after_line = if *start_line > 0 { Some(start_line - 1) } else { None };
                editor.insert_markdown_cell(content, after_line)?;
            }
            UndoableOperation::EditMarkdownCell { start_line, end_line, new_content, is_module_doc, .. } => {
                // Redo edit markdown = apply new content again
                editor.edit_markdown_cell(*start_line, *end_line, new_content, *is_module_doc)?;
            }
            UndoableOperation::DeleteMarkdownCell { start_line, content } => {
                // Redo delete markdown = delete again
                // We need to find the end line by counting content lines
                let line_count = content.lines().count();
                let end_line = start_line + line_count;
                editor.delete_markdown_cell(*start_line, end_line)?;
            }
            UndoableOperation::MoveMarkdownCell { start_line, end_line, direction } => {
                // Redo move markdown = move in same direction
                editor.move_markdown_cell(*start_line, *end_line, *direction)?;
            }
        }

        editor.save()?;

        // Record for undo (so we can undo the redo)
        self.undo_manager.record(operation);

        // Reload to update in-memory state
        self.reload()?;

        Ok(description)
    }

    /// Get the current undo/redo state.
    pub fn get_undo_redo_state(&self) -> ServerMessage {
        ServerMessage::UndoRedoState {
            can_undo: self.undo_manager.can_undo(),
            can_redo: self.undo_manager.can_redo(),
            undo_description: self.undo_manager.undo_description(),
            redo_description: self.undo_manager.redo_description(),
        }
    }

    /// Clear undo/redo history.
    ///
    /// Called when the file is externally modified.
    pub fn clear_undo_history(&mut self) {
        self.undo_manager.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_session_creation() {
        // This would require a real notebook file, so we just test the types compile
        let (tx, _rx) = broadcast::channel::<ServerMessage>(16);
        drop(tx);
    }
}