From 303f816e76bdc161221bce55cfd7fd054ddc4749 Mon Sep 17 00:00:00 2001 From: dawsh2 <104176151+dawsh2@users.noreply.github.com> Date: Wed, 5 Nov 2025 20:34:15 -0800 Subject: [PATCH] Experiment with inline Mermaid preview --- lsp/src/main.rs | 722 +++++++++------------------------- lsp/src/render.rs | 23 +- lsp/tests/integration_test.rs | 85 ++-- 3 files changed, 258 insertions(+), 572 deletions(-) diff --git a/lsp/src/main.rs b/lsp/src/main.rs index 36a8d2a..4c91c36 100644 --- a/lsp/src/main.rs +++ b/lsp/src/main.rs @@ -3,31 +3,31 @@ use log::{debug, error, info, warn}; use lsp_server::{Connection, Message, Request, RequestId, Response, ResponseError}; use lsp_types::*; use serde_json::json; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; use std::{ collections::HashMap, fs, path::Path, - sync::atomic::{AtomicUsize, Ordering}, + sync::atomic::{AtomicU64, AtomicUsize, Ordering}, time::{SystemTime, UNIX_EPOCH}, }; use url::Url; -use std::collections::hash_map::DefaultHasher; -use std::hash::{Hash, Hasher}; mod render; use crate::render::render_mermaid; // Constants to avoid magic strings -const MERMAID_SOURCE_COMMENT_PREFIX: &str = ""; const MERMAID_MEDIA_DIR: &str = ".mermaid"; const MERMAID_CACHE_DIR: &str = ".cache"; -const MERMAID_FILE_EXTENSION: &str = ".mmd"; const MERMAID_FENCE_START: &str = "```mermaid"; -const MERMAID_FENCE_END: &str = "```"; +const MERMAID_PREVIEW_COMMENT_PREFIX: &str = ""; +const MERMAID_SOURCE_SUMMARY: &str = "Show Mermaid source"; static SVG_COUNTER: AtomicUsize = AtomicUsize::new(0); +static REQUEST_ID_COUNTER: AtomicU64 = AtomicU64::new(1); /// Send an error notification to the LSP client fn send_error_notification(connection: &Connection, message: &str) { @@ -60,90 +60,6 @@ fn send_warning_notification(connection: &Connection, message: &str) { } } -// Strip mermaid wrapper (```mermaid ... ```) from code if present -fn strip_mermaid_wrapper(code: &str) -> String { - let trimmed = code.trim(); - let lines: Vec<&str> = trimmed.lines().collect(); - - if lines.is_empty() { - return code.to_string(); - } - - let has_start = lines[0].trim().starts_with(MERMAID_FENCE_START); - let has_end = lines.last().map(|l| l.trim() == MERMAID_FENCE_END).unwrap_or(false); - - if has_start && has_end && lines.len() >= 2 { - return lines[1..lines.len() - 1].join("\n"); - } - - code.to_string() -} - -// Find the most recent matching source file when the referenced file doesn't exist -fn find_most_recent_source_file(missing_path: &Path, _uri: &str) -> Option { - debug!("Searching for recent source file matching pattern"); - - // Extract the base filename pattern from the missing path - if let Some(file_name) = missing_path.file_name().and_then(|n| n.to_str()) { - // Extract base name and diagram number (e.g., "example_0" from "example_1761843815_0.mmd") - let parts: Vec<&str> = file_name.split('_').collect(); - if parts.len() >= 3 { - let base_name = parts[0]; // e.g., "example" - let diagram_num = parts[parts.len() - 2]; // e.g., "0" - let extension = parts[parts.len() - 1]; // e.g., "mmd" - - // Construct search pattern - let pattern = format!("{}_{}_{}", base_name, "*", diagram_num); - - // Get the directory to search in - let search_dir = missing_path.parent().unwrap_or_else(|| Path::new(MERMAID_MEDIA_DIR)); - - debug!("Searching in {:?} for pattern {}", search_dir, pattern); - - // Find all matching files and get the most recent one - if let Ok(entries) = std::fs::read_dir(search_dir) { - let mut best_match: Option<(std::fs::DirEntry, std::time::SystemTime)> = None; - - for entry in entries.flatten() { - let path = entry.path(); - if let Some(name) = path.file_name().and_then(|n| n.to_str()) { - // Check if it matches our pattern - if name.starts_with(&format!("{}_{}", base_name, diagram_num)) && name.ends_with(&format!(".{}", extension)) { - // Get modification time - if let Ok(metadata) = entry.metadata() { - if let Ok(modified) = metadata.modified() { - match &best_match { - None => best_match = Some((entry, modified)), - Some((_, best_time)) => { - if modified > *best_time { - best_match = Some((entry, modified)); - } - } - } - } - } - } - } - } - - if let Some((best_entry, _)) = best_match { - let best_path = best_entry.path(); - debug!("Found most recent match: {:?}", best_path); - - // Try to read it - if let Ok(content) = std::fs::read_to_string(&best_path) { - debug!("Successfully read recent file ({} bytes)", content.len()); - return Some(content); - } - } - } - } - } - - debug!("No recent source file found"); - None -} - fn main() -> Result<()> { // Initialize logging to a file so we can actually see what's happening let log_file = Path::new("/tmp/mermaid-lsp.log"); @@ -152,8 +68,8 @@ fn main() -> Result<()> { if let Ok(file) = std::fs::OpenOptions::new() .create(true) .append(true) - .open(&log_file) { - + .open(&log_file) + { env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")) .format_timestamp_millis() .target(env_logger::Target::Pipe(Box::new(file))) @@ -277,7 +193,7 @@ fn handle_request( info!("URI: {}", params.text_document.uri); info!("Range: {:?}", params.range); - let actions = get_code_actions(¶ms, documents, connection)?; + let actions = get_code_actions(¶ms, documents)?; info!("Returning {} code actions", actions.len()); for action in &actions { @@ -390,7 +306,6 @@ fn handle_notification( fn get_code_actions( params: &CodeActionParams, documents: &HashMap, - _connection: &Connection, ) -> Result> { let uri = params.text_document.uri.to_string(); let cursor = params.range.start; @@ -409,72 +324,53 @@ fn get_code_actions( // Count total mermaid blocks in the document - O(1) operation let total_blocks = count_mermaid_blocks(content); - info!("Found {} mermaid blocks, cursor at line {}", total_blocks, cursor.line); + info!( + "Found {} mermaid blocks, cursor at line {}", + total_blocks, cursor.line + ); - // Render All - pre-compute edit for Zed compatibility if total_blocks > 1 { - info!("Adding Render All action for {} diagrams (pre-computing edit)", total_blocks); - - // Pre-compute the WorkspaceEdit - info!("Calling render_all_diagrams_content..."); - match render_all_diagrams_content(&uri, content, Some(_connection)) { - Ok(changes) => { - info!("Successfully rendered all diagrams, got {} file changes", changes.len()); - let edit = WorkspaceEdit { - changes: Some(changes), - document_changes: None, - change_annotations: None, - }; - - actions.push(CodeAction { - title: format!("Render All {} Mermaid Diagrams", total_blocks), - kind: Some(CodeActionKind::REFACTOR_REWRITE), - diagnostics: None, - edit: Some(edit), // Direct edit, no command - command: None, - is_preferred: Some(true), - disabled: None, - data: None, - }); - info!("Render All action added successfully"); - } - Err(e) => { - error!("Failed to pre-compute Render All edit: {}", e); - } - } + info!("Adding Render All action for {} diagrams", total_blocks); + let arguments = vec![json!({ "uri": uri })]; + actions.push(CodeAction { + title: format!("Render All {} Mermaid Diagrams", total_blocks), + kind: Some(CodeActionKind::REFACTOR_REWRITE), + diagnostics: None, + edit: None, + command: Some(Command { + title: "Render All Mermaid Diagrams".to_string(), + command: "mermaid.renderAllLightweight".to_string(), + arguments: Some(arguments), + }), + is_preferred: Some(true), + disabled: None, + data: None, + }); } else { info!("Not adding Render All (only {} blocks)", total_blocks); } - // Edit All - pre-compute edit for Zed compatibility let rendered_count = count_rendered_blocks(content); if rendered_count > 1 { - debug!("Adding Edit All action for {} rendered diagrams (pre-computing edit)", rendered_count); - - // Pre-compute the WorkspaceEdit - match edit_all_sources_content(&uri, content) { - Ok(changes) => { - let edit = WorkspaceEdit { - changes: Some(changes), - document_changes: None, - change_annotations: None, - }; - - actions.push(CodeAction { - title: format!("Edit All {} Mermaid Sources", rendered_count), - kind: Some(CodeActionKind::REFACTOR_REWRITE), - diagnostics: None, - edit: Some(edit), // Direct edit, no command - command: None, - is_preferred: Some(false), - disabled: None, - data: None, - }); - } - Err(e) => { - warn!("Failed to pre-compute Edit All edit: {}", e); - } - } + debug!( + "Adding Edit All action for {} rendered diagrams", + rendered_count + ); + let arguments = vec![json!({ "uri": uri })]; + actions.push(CodeAction { + title: format!("Edit All {} Mermaid Sources", rendered_count), + kind: Some(CodeActionKind::REFACTOR_REWRITE), + diagnostics: None, + edit: None, + command: Some(Command { + title: "Edit All Mermaid Sources".to_string(), + command: "mermaid.editAllSources".to_string(), + arguments: Some(arguments), + }), + is_preferred: Some(false), + disabled: None, + data: None, + }); } // Render Single - skip for now, only support bulk operations @@ -489,9 +385,12 @@ fn get_code_actions( if cursor_line < lines.len() { let line = lines[cursor_line].trim(); - let is_on_comment = line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) && line.ends_with(MERMAID_SOURCE_COMMENT_SUFFIX); + let is_on_comment = line == MERMAID_INLINE_SOURCE_COMMENT; - debug!("Line {}: '{}' - is_comment: {}", cursor_line, line, is_on_comment); + debug!( + "Line {}: '{}' - is_comment: {}", + cursor_line, line, is_on_comment + ); // Skip Edit Single for now - only support Edit All debug!("Cursor state checked, skipping Edit Single action"); @@ -530,219 +429,86 @@ fn is_mermaid_document(uri: &str) -> bool { uri.ends_with(".mmd") || uri.ends_with(".mermaid") } -fn locate_mermaid_source_block( - content: &str, - uri: &str, - cursor: &Position, -) -> Option { - if is_mermaid_document(uri) { - let lines: Vec<&str> = content.lines().collect(); - let last_line = lines.len().saturating_sub(1); - let end_character = lines.get(last_line).map(|l| l.len()).unwrap_or(0); - - // Strip mermaid wrapper if present - let clean_code = strip_mermaid_wrapper(content); - - return Some(MermaidSourceBlock { - code: clean_code, - start: Position { - line: 0, - character: 0, - }, - end: Position { - line: last_line as u32, - character: end_character as u32, - }, - kind: DocumentKind::Mermaid, - }); - } - - let lines: Vec<&str> = content.lines().collect(); - if lines.is_empty() { - return None; - } - - let cursor_line = cursor.line.min((lines.len() - 1) as u32) as usize; - let (start_line, end_line) = find_mermaid_fence(&lines, cursor_line)?; - - // Check if this block is already rendered (has source file comment before it) - if start_line > 0 { - let prev_line = lines[start_line - 1].trim(); - if prev_line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) { - return None; - } - } - - let code = lines[start_line + 1..end_line].join("\n"); - - let end_position = if end_line + 1 < lines.len() { - Position { - line: (end_line + 1) as u32, - character: 0, - } - } else { - Position { - line: end_line as u32, - character: lines[end_line].len() as u32, - } - }; - - Some(MermaidSourceBlock { - code, - start: Position { - line: start_line as u32, - character: 0, - }, - end: end_position, - kind: DocumentKind::Markdown, - }) -} - fn locate_rendered_mermaid_block( content: &str, uri: &str, cursor: &Position, ) -> Option { - debug!("locate_rendered_mermaid_block ENTRY - content length: {}", content.len()); - let lines: Vec<&str> = content.lines().collect(); - debug!("locate_rendered_mermaid_block - parsed {} lines", lines.len()); - if lines.is_empty() { - debug!("locate_rendered_mermaid_block - EARLY RETURN: lines.is_empty()"); return None; } let cursor_line = cursor.line.min((lines.len() - 1) as u32) as usize; - debug!("locate_rendered_mermaid_block - cursor at line {}, total lines: {}", cursor_line, lines.len()); - - debug!("=== locate_rendered_mermaid_block called ==="); - debug!("Cursor line: {}, total lines: {}", cursor_line, lines.len()); - - // Find comment with mermaid source file reference - // Search BACKWARDS from cursor first (most common: cursor on image line after comment) - // Then search forward if not found - debug!("Searching for mermaid comment near cursor line {}", cursor_line); - - let source_line = { - // First, search backwards from cursor (up to 10 lines) - let search_start = cursor_line.saturating_sub(10); - let backward_result = (search_start..=cursor_line).rev().find(|&i| { - let line = lines[i].trim(); - let is_comment = line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) && line.ends_with(MERMAID_SOURCE_COMMENT_SUFFIX); - if is_comment { - debug!("Found mermaid comment (backward) at line {}: {}", i, line); - } - is_comment - }); - if let Some(line) = backward_result { - line - } else { - // If not found backward, search forward (up to 5 lines) - let search_end = (cursor_line + 5).min(lines.len() - 1); - let forward_result = (cursor_line..=search_end).find(|&i| { - let line = lines[i].trim(); - let is_comment = line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) && line.ends_with(MERMAID_SOURCE_COMMENT_SUFFIX); - if is_comment { - debug!("Found mermaid comment (forward) at line {}: {}", i, line); - } - is_comment - }); + // Locate the preview comment that anchors a rendered block + let preview_line = { + let search_start = cursor_line.saturating_sub(15); + let backward = (search_start..=cursor_line) + .rev() + .find(|&i| lines[i].trim().starts_with(MERMAID_PREVIEW_COMMENT_PREFIX)); - forward_result? - } - }; - - // Extract the source file path - let line = lines[source_line].trim(); - let file_start = MERMAID_SOURCE_COMMENT_PREFIX.len(); - let file_end = line.len() - "-->".len(); - let source_file_path = &line[file_start..file_end].trim(); - - // Get the full path to the source file - let source_full_path = if let Ok(url) = Url::parse(uri) { - if let Some(path) = url.to_file_path().ok() { - // source_file_path is relative to the document's parent - if let Some(parent) = path.parent() { - let full_path = parent.join(source_file_path); - debug!("Document path: {:?}", path); - debug!("Document parent: {:?}", parent); - debug!("Relative source file: {}", source_file_path); - debug!("Resolved full path: {:?}", full_path); - - full_path - } else { - debug!("No parent directory for document, using relative path"); - Path::new(source_file_path).to_path_buf() - } + if let Some(idx) = backward { + Some(idx) } else { - debug!("Could not parse URI to file path: {}", uri); - Path::new(source_file_path).to_path_buf() + let search_end = (cursor_line + 15).min(lines.len().saturating_sub(1)); + (cursor_line..=search_end) + .find(|&i| lines[i].trim().starts_with(MERMAID_PREVIEW_COMMENT_PREFIX)) } - } else { - debug!("Could not parse URI: {}", uri); - Path::new(source_file_path).to_path_buf() - }; - - debug!("Looking for source file at: {:?}", source_full_path); - debug!("File exists: {}", source_full_path.exists()); - - // Read the source from the file - let code = match fs::read_to_string(&source_full_path) { - Ok(content) => { - debug!("Successfully read source file ({} bytes)", content.len()); - content + }?; + + // Find the inline source marker and fenced code block that follows it + let mut inline_comment_line = None; + for idx in preview_line + 1..lines.len() { + let trimmed = lines[idx].trim(); + if trimmed == MERMAID_INLINE_SOURCE_COMMENT { + inline_comment_line = Some(idx); + break; } - Err(e) => { - debug!("Failed to read source file: {}, attempting to find recent file...", e); - debug!("Error details: {:?}", e.kind()); - - // Try to find the most recent matching file - if let Some(recent_code) = find_most_recent_source_file(&source_full_path, &uri) { - debug!("Found recent source file, using that instead"); - recent_code - } else { - debug!("Could not find any recent source file"); - return None; - } + if trimmed.starts_with(MERMAID_PREVIEW_COMMENT_PREFIX) { + break; } - }; + } + let inline_comment_line = inline_comment_line?; - // Find the image reference (usually on the next non-empty line) - let mut img_line = source_line + 1; - while img_line < lines.len() && lines[img_line].trim().is_empty() { - img_line += 1; + let code_start_line = inline_comment_line + 1; + if code_start_line >= lines.len() { + return None; } - // Find the end of the block (after the image and any trailing blank lines) - let end_line = if img_line < lines.len() && lines[img_line].contains("![Mermaid Diagram](") { - // Start after the image line - let mut end = img_line + 1; + if lines[code_start_line].trim_start() != MERMAID_FENCE_START { + return None; + } - // Skip ONE blank line if present (common formatting) - if end < lines.len() && lines[end].trim().is_empty() { - end += 1; + let mut code_end_line = None; + for idx in code_start_line + 1..lines.len() { + if lines[idx].trim_start().starts_with("```") { + code_end_line = Some(idx); + break; } + } + let code_end_line = code_end_line?; - // But stop if we hit another diagram or content - // Don't consume the next diagram's comment or headers - end - } else { - source_line + 2 - }; - - debug!("Found rendered block - comment line {}, img line {}, end line {}", source_line, img_line, end_line); + let code = lines[code_start_line + 1..code_end_line].join("\n"); + // Find the closing + let mut details_end_line = None; + for idx in code_end_line + 1..lines.len() { + if lines[idx].trim().starts_with("") { + details_end_line = Some(idx + 1); + break; + } + } + let details_end_line = details_end_line.unwrap_or(code_end_line + 1); Some(RenderedMermaidBlock { code, start: Position { - line: source_line as u32, + line: preview_line as u32, character: 0, }, end: Position { - line: end_line.min(lines.len()) as u32, + line: details_end_line as u32, character: 0, }, kind: if is_mermaid_document(uri) { @@ -776,79 +542,6 @@ fn find_mermaid_fence(lines: &[&str], cursor_line: usize) -> Option<(usize, usiz Some((start, end)) } -/// Clean up old diagram files that are no longer referenced in the document -/// Keeps cache files (.cache/*.svg) but removes unreferenced output files -fn cleanup_old_diagram_files(_uri: &str, content: &str, media_dir: &Path) -> Result<()> { - info!("=== CLEANUP: Cleaning up old diagram files in {:?}", media_dir); - - // Find all currently referenced files in the document - let mut referenced_files = std::collections::HashSet::new(); - for line in content.lines() { - if line.contains(MERMAID_SOURCE_COMMENT_PREFIX) { - // Extract the .mmd file path from comment - if let Some(start) = line.find(MERMAID_SOURCE_COMMENT_PREFIX) { - let path_start = start + MERMAID_SOURCE_COMMENT_PREFIX.len(); - if let Some(end) = line[path_start..].find(MERMAID_SOURCE_COMMENT_SUFFIX) { - let file_path = line[path_start..path_start + end].trim(); - referenced_files.insert(file_path.to_string()); - } - } - } - // Also collect SVG references from markdown image links - if line.contains("![Mermaid Diagram](") { - if let Some(start) = line.find("](") { - let path_start = start + 2; - if let Some(end) = line[path_start..].find(')') { - let file_path = line[path_start..path_start + end].trim(); - referenced_files.insert(file_path.to_string()); - } - } - } - } - - info!("CLEANUP: Found {} referenced files in document", referenced_files.len()); - for ref_file in &referenced_files { - info!("CLEANUP: Referenced: {}", ref_file); - } - - // Scan the media directory for orphaned files - if let Ok(entries) = std::fs::read_dir(media_dir) { - for entry in entries.flatten() { - let path = entry.path(); - - // Skip directories (like .cache) - if path.is_dir() { - continue; - } - - // Only clean up .mmd and .svg files - if let Some(ext) = path.extension() { - if ext != "mmd" && ext != "svg" { - continue; - } - } else { - continue; - } - - // Check if this file is referenced - let file_name = path.file_name().unwrap().to_string_lossy(); - let relative_path = format!("{}/{}", MERMAID_MEDIA_DIR, file_name); - - if !referenced_files.contains(file_name.as_ref()) && - !referenced_files.contains(&relative_path) { - info!("CLEANUP: Removing unreferenced file: {:?}", path); - if let Err(e) = std::fs::remove_file(&path) { - warn!("CLEANUP: Failed to remove old file {:?}: {}", path, e); - } - } else { - info!("CLEANUP: Keeping referenced file: {:?}", path); - } - } - } - - Ok(()) -} - fn create_render_edits( uri: &str, block: &MermaidSourceBlock, @@ -885,7 +578,9 @@ fn create_render_edits( } } } else { - return Err(anyhow!("Path traversal attempt detected: path contains '..'")); + return Err(anyhow!( + "Path traversal attempt detected: path contains '..'" + )); } } } @@ -914,8 +609,7 @@ fn create_render_edits( // Check if we have a cached version let svg_contents = if cache_path.exists() { debug!("Using cached SVG for hash {:x}", code_hash); - fs::read_to_string(&cache_path) - .map_err(|e| anyhow!("Failed to read cached SVG: {}", e))? + fs::read_to_string(&cache_path).map_err(|e| anyhow!("Failed to read cached SVG: {}", e))? } else { debug!("Rendering new SVG (cache miss) for hash {:x}", code_hash); let contents = render_mermaid(&block.code)?; @@ -951,38 +645,28 @@ fn create_render_edits( .map_err(|e| anyhow!("Failed to write SVG: {}", e))?; info!("Successfully wrote SVG file"); - let source_file_path = { - let base_name = path.file_stem() - .unwrap_or_default() - .to_string_lossy(); - let source_filename = format!("{}_{}{}", base_name, unique_id, MERMAID_FILE_EXTENSION); - media_dir.join(source_filename) - }; - - // Write the source to the .mmd file - fs::write(&source_file_path, &block.code) - .map_err(|e| anyhow!("Failed to write source file: {}", e))?; - - // Calculate relative paths from the markdown file to mermaid media directory - let source_relative = source_file_path - .strip_prefix(&path.parent().unwrap_or_else(|| Path::new("."))) - .unwrap_or(&source_file_path) - .to_string_lossy(); - let svg_path_buf = Path::new(MERMAID_MEDIA_DIR).join(&svg_filename); let svg_relative = svg_path_buf.to_string_lossy(); - let mut new_text = format!( - "{}{}{}\n\n![Mermaid Diagram]({})\n", - MERMAID_SOURCE_COMMENT_PREFIX, source_relative, MERMAID_SOURCE_COMMENT_SUFFIX, - svg_relative - ); - - debug!("Rendering with external source file"); - - if !new_text.ends_with('\n') { - new_text.push('\n'); - } + let preview_comment = format!("{}{} -->", MERMAID_PREVIEW_COMMENT_PREFIX, svg_relative); + + let mut new_text = String::new(); + new_text.push_str(&preview_comment); + new_text.push('\n'); + new_text.push_str("
\n"); + new_text.push_str(&format!("![Mermaid Diagram]({})\n", svg_relative)); + new_text.push_str("
\n\n"); + new_text.push_str("
\n"); + new_text.push_str(&format!( + " {}\n", + MERMAID_SOURCE_SUMMARY + )); + new_text.push_str(&format!(" {}\n", MERMAID_INLINE_SOURCE_COMMENT)); + new_text.push_str("```mermaid\n"); + new_text.push_str(block.code.trim_end()); + new_text.push('\n'); + new_text.push_str("```\n"); + new_text.push_str("
\n"); let mut changes = HashMap::new(); changes.insert( @@ -1044,7 +728,7 @@ fn count_mermaid_blocks(content: &str) -> usize { while i < lines.len() { if let Some((start, end)) = find_mermaid_fence(&lines, i) { // Check if it's already rendered - if start == 0 || !lines[start - 1].starts_with(MERMAID_SOURCE_COMMENT_PREFIX) { + if start == 0 || lines[start - 1].trim() != MERMAID_INLINE_SOURCE_COMMENT { count += 1; } i = end + 1; @@ -1061,7 +745,7 @@ fn count_rendered_blocks(content: &str) -> usize { let mut count = 0; for line in lines { - if line.trim().starts_with(MERMAID_SOURCE_COMMENT_PREFIX) { + if line.trim().starts_with(MERMAID_PREVIEW_COMMENT_PREFIX) { count += 1; } } @@ -1069,10 +753,7 @@ fn count_rendered_blocks(content: &str) -> usize { count } -fn edit_all_sources_content( - uri: &str, - content: &str, -) -> Result>> { +fn edit_all_sources_content(uri: &str, content: &str) -> Result>> { let lines: Vec<&str> = content.lines().collect(); let mut all_edits: HashMap> = HashMap::new(); let mut i = 0; @@ -1082,73 +763,42 @@ fn edit_all_sources_content( while i < lines.len() { let line = lines[i].trim(); - // Look for mermaid source comment lines - if line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) && line.ends_with(MERMAID_SOURCE_COMMENT_SUFFIX) { + if line.starts_with(MERMAID_PREVIEW_COMMENT_PREFIX) { debug!("Found rendered block at line {}", i); - // Find the end of the rendered block (next blank line or mermaid fence) - let mut end = i + 1; - while end < lines.len() { - let next_line = lines[end].trim(); - if next_line.is_empty() || next_line.starts_with("```mermaid") || next_line.starts_with(MERMAID_SOURCE_COMMENT_PREFIX) { - break; - } - end += 1; - } + let cursor = Position { + line: i as u32, + character: 0, + }; - // Extract the source file path from comment - let start_pos = line.find(MERMAID_SOURCE_COMMENT_PREFIX).unwrap() + MERMAID_SOURCE_COMMENT_PREFIX.len(); - let end_pos = line.len() - MERMAID_SOURCE_COMMENT_SUFFIX.len(); - let source_file = &line[start_pos..end_pos]; - - debug!("Loading source from: {}", source_file); - - // Read the source file - if let Ok(source_url) = Url::parse(uri) { - if let Ok(doc_path) = source_url.to_file_path() { - if let Some(parent) = doc_path.parent() { - let source_path = parent.join(source_file); - if let Ok(source_code) = std::fs::read_to_string(&source_path) { - // Create the block - let block = RenderedMermaidBlock { - code: source_code, - start: Position { - line: i as u32, - character: 0, - }, - end: Position { - line: end as u32, - character: 0, - }, - kind: DocumentKind::Markdown, - }; - - match create_source_edits(uri, &block) { - Ok(mut edits) => { - if let Some((url, mut text_edits)) = edits.drain().next() { - if let Some(existing_edits) = all_edits.get_mut(&url) { - existing_edits.append(&mut text_edits); - } else { - all_edits.insert(url, text_edits); - } - } - } - Err(e) => { - warn!("Failed to create source edits for line {}: {}", i + 1, e); - } + if let Some(block) = locate_rendered_mermaid_block(content, uri, &cursor) { + match create_source_edits(uri, &block) { + Ok(mut edits) => { + if let Some((url, mut text_edits)) = edits.drain().next() { + if let Some(existing_edits) = all_edits.get_mut(&url) { + existing_edits.append(&mut text_edits); + } else { + all_edits.insert(url, text_edits); } } } + Err(e) => { + warn!("Failed to create source edits for line {}: {}", i + 1, e); + } } - } - i = end; - } else { - i += 1; + i = block.end.line as usize; + continue; + } } + + i += 1; } - debug!("Found {} sets of edits across all rendered blocks", all_edits.len()); + debug!( + "Found {} sets of edits across all rendered blocks", + all_edits.len() + ); Ok(all_edits) } @@ -1159,14 +809,14 @@ fn render_all_diagrams_content( ) -> Result>> { let lines: Vec<&str> = content.lines().collect(); let mut all_edits: HashMap> = HashMap::new(); - let mut rendered_any = false; // Track if we actually rendered anything + let mut rendered_any = false; // Track if we actually rendered anything let mut i = 0; while i < lines.len() { if let Some((start, end)) = find_mermaid_fence(&lines, i) { // Skip if already rendered - if start == 0 || !lines[start - 1].starts_with(MERMAID_SOURCE_COMMENT_PREFIX) { - rendered_any = true; // Mark that we're rendering something + if start == 0 || lines[start - 1].trim() != MERMAID_INLINE_SOURCE_COMMENT { + rendered_any = true; // Mark that we're rendering something let code = lines[start + 1..end].join("\n"); let block = MermaidSourceBlock { @@ -1204,7 +854,8 @@ fn render_all_diagrams_content( } } Err(e) => { - let error_msg = format!("Failed to render diagram at line {}: {}", start + 1, e); + let error_msg = + format!("Failed to render diagram at line {}: {}", start + 1, e); error!("{}", error_msg); if let Some(conn) = connection { send_error_notification(conn, &error_msg); @@ -1231,11 +882,7 @@ fn render_all_diagrams_content( Ok(all_edits) } -fn apply_workspace_edit( - connection: &Connection, - edit: WorkspaceEdit, - label: &str, -) -> Result<()> { +fn apply_workspace_edit(connection: &Connection, edit: WorkspaceEdit, label: &str) -> Result<()> { info!("Sending workspace/applyEdit request: {}", label); let params = ApplyWorkspaceEditParams { @@ -1243,10 +890,11 @@ fn apply_workspace_edit( edit, }; + let request_id = REQUEST_ID_COUNTER.fetch_add(1, Ordering::Relaxed); let request = Request::new( - RequestId::from(std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?.as_nanos() as i32), + RequestId::from(request_id.to_string()), "workspace/applyEdit".to_string(), - serde_json::to_value(params)? + serde_json::to_value(params)?, ); connection.sender.send(Message::Request(request))?; @@ -1265,7 +913,8 @@ fn execute_command( match params.command.as_str() { "mermaid.renderAllLightweight" => { // Get URI from command arguments - let uri = params.arguments + let uri = params + .arguments .first() .and_then(|arg| arg.get("uri")) .and_then(|v| v.as_str()) @@ -1290,7 +939,8 @@ fn execute_command( } "mermaid.renderSingle" => { // Get parameters from command arguments - let args = params.arguments + let args = params + .arguments .first() .ok_or_else(|| anyhow::anyhow!("No arguments provided"))?; @@ -1299,15 +949,15 @@ fn execute_command( .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing URI argument"))?; - let start_line = args - .get("startLine") - .and_then(|v| v.as_u64()) - .ok_or_else(|| anyhow::anyhow!("Missing startLine"))? as u32; + let start_line = + args.get("startLine") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("Missing startLine"))? as u32; - let end_line = args - .get("endLine") - .and_then(|v| v.as_u64()) - .ok_or_else(|| anyhow::anyhow!("Missing endLine"))? as u32; + let end_line = + args.get("endLine") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("Missing endLine"))? as u32; let code = args .get("code") @@ -1343,7 +993,8 @@ fn execute_command( Ok(()) } "mermaid.editSingleSource" => { - let args = params.arguments + let args = params + .arguments .first() .ok_or_else(|| anyhow::anyhow!("No arguments provided"))?; @@ -1352,15 +1003,15 @@ fn execute_command( .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing URI argument"))?; - let start_line = args - .get("startLine") - .and_then(|v| v.as_u64()) - .ok_or_else(|| anyhow::anyhow!("Missing startLine"))? as u32; + let start_line = + args.get("startLine") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("Missing startLine"))? as u32; - let end_line = args - .get("endLine") - .and_then(|v| v.as_u64()) - .ok_or_else(|| anyhow::anyhow!("Missing endLine"))? as u32; + let end_line = + args.get("endLine") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("Missing endLine"))? as u32; let code = args .get("code") @@ -1394,7 +1045,8 @@ fn execute_command( Ok(()) } "mermaid.editAllSources" => { - let uri = params.arguments + let uri = params + .arguments .first() .and_then(|arg| arg.get("uri")) .and_then(|v| v.as_str()) diff --git a/lsp/src/render.rs b/lsp/src/render.rs index 33c01c6..b2b95c2 100644 --- a/lsp/src/render.rs +++ b/lsp/src/render.rs @@ -1,13 +1,13 @@ use anyhow::{anyhow, Result}; +use html_escape; +use once_cell::sync::Lazy; +use regex::Regex; use std::{ env, fs, path::PathBuf, process::{Command, Stdio}, }; use tempfile::tempdir; -use once_cell::sync::Lazy; -use regex::Regex; -use html_escape; // Precompiled regex patterns to avoid DoS and improve performance static FOREIGN_OBJECT_REGEX: Lazy = Lazy::new(|| { @@ -173,8 +173,7 @@ fn extract_text_from_html(html: &str) -> String { } fn extract_attr(tag: &str, attr: &str) -> Option { - let attr_regex = Regex::new(&format!(r#"{}="([^"]*)""#, regex::escape(attr))) - .ok()?; + let attr_regex = Regex::new(&format!(r#"{}="([^"]*)""#, regex::escape(attr))).ok()?; attr_regex.captures(tag).map(|c| c[1].to_string()) } @@ -210,9 +209,8 @@ static JAVASCRIPT_HREF_ATTR: Lazy = Lazy::new(|| { .expect("valid regex for javascript href attributes") }); -static HTML_TAG_REGEX: Lazy = Lazy::new(|| { - Regex::new(r"<[^>]*>").expect("valid regex for HTML tags") -}); +static HTML_TAG_REGEX: Lazy = + Lazy::new(|| Regex::new(r"<[^>]*>").expect("valid regex for HTML tags")); #[cfg(test)] mod tests { @@ -235,7 +233,10 @@ mod tests { ]; for svg in test_cases { - assert!(sanitize_svg(svg).is_err(), "Should reject case-insensitive script tags"); + assert!( + sanitize_svg(svg).is_err(), + "Should reject case-insensitive script tags" + ); } } @@ -331,9 +332,9 @@ mod tests { assert!(!sanitized.contains("onclick=\"alert('xss')\"")); assert!(!sanitized.contains("alert('xss')\"")); assert!(!sanitized.contains("…")); // ellipsis from truncation - // Should be well-formed + // Should be well-formed assert!(sanitized.contains("")); } -} \ No newline at end of file +} diff --git a/lsp/tests/integration_test.rs b/lsp/tests/integration_test.rs index dc13aeb..5a5707d 100644 --- a/lsp/tests/integration_test.rs +++ b/lsp/tests/integration_test.rs @@ -38,7 +38,11 @@ fn test_path_traversal_patterns() { ]; for path in malicious_paths { - assert!(path.contains(".."), "Path should contain '..' pattern: {}", path); + assert!( + path.contains(".."), + "Path should contain '..' pattern: {}", + path + ); } } @@ -47,21 +51,25 @@ fn test_path_traversal_patterns() { fn test_safe_paths() { let safe_paths = vec![ ".mermaid/diagram.svg", - ".mermaid/example_123.mmd", + ".mermaid/example_123.svg", "subfolder/diagram.svg", ]; for path in safe_paths { - assert!(!path.contains(".."), "Safe path should not contain '..': {}", path); + assert!( + !path.contains(".."), + "Safe path should not contain '..': {}", + path + ); } } /// Test file extension validation #[test] fn test_file_extension_validation() { - let valid_extensions = vec![".mmd", ".svg", ".md"]; + let valid_extensions = vec![".svg", ".md"]; let test_files = vec![ - ("diagram.mmd", true), + ("diagram.mmd", false), ("diagram.svg", true), ("document.md", true), ("script.sh", false), @@ -76,7 +84,11 @@ fn test_file_extension_validation() { .unwrap_or_default(); let is_valid = valid_extensions.contains(&ext.as_str()); - assert_eq!(is_valid, should_be_valid, "File '{}' validation mismatch", filename); + assert_eq!( + is_valid, should_be_valid, + "File '{}' validation mismatch", + filename + ); } } @@ -125,29 +137,31 @@ sequenceDiagram fence_count += 1; // Verify closing fence exists let closing_fence = lines.iter().skip(i + 1).position(|l| l.trim() == "```"); - assert!(closing_fence.is_some(), "Mermaid fence should have closing fence"); + assert!( + closing_fence.is_some(), + "Mermaid fence should have closing fence" + ); } } assert_eq!(fence_count, 2, "Should find 2 mermaid fences"); } -/// Test source comment format +/// Test preview comment format #[test] fn test_source_comment_format() { - let comment = ""; + let comment = ""; - assert!(comment.starts_with("")); - // Extract path - let start = "".len(); - let path = &comment[start..end]; + let path = comment[start..end].trim(); - assert_eq!(path, ".mermaid/example_123.mmd"); + assert_eq!(path, ".mermaid/example_123.svg"); assert!(path.starts_with(".mermaid/")); - assert!(path.ends_with(".mmd")); + assert!(path.ends_with(".svg")); } /// Test SVG validation - reject scripts @@ -156,8 +170,14 @@ fn test_svg_script_rejection() { let malicious_svg = r#""#; let safe_svg = r#""#; - assert!(malicious_svg.contains("]+>([^<]+(?:<(?!/foreignObject>)[^<]*)*)"#; // Verify pattern structure doesn't have dangerous patterns - assert!(!pattern.contains(".*?)*"), "Should not have nested greedy quantifiers"); - assert!(!pattern.contains(".+)+"), "Should not have nested possessive quantifiers"); + assert!( + !pattern.contains(".*?)*"), + "Should not have nested greedy quantifiers" + ); + assert!( + !pattern.contains(".+)+"), + "Should not have nested possessive quantifiers" + ); // The pattern uses [^<]+ and [^>]+ which are safe because they're negated character classes - assert!(pattern.contains("[^<]"), "Should use negated character classes"); - assert!(pattern.contains("[^>]"), "Should use negated character classes"); + assert!( + pattern.contains("[^<]"), + "Should use negated character classes" + ); + assert!( + pattern.contains("[^>]"), + "Should use negated character classes" + ); } /// Test cleanup file detection @@ -218,12 +254,10 @@ fn test_cleanup_file_detection() { // Create some test files let old_svg = media_dir.join("old_diagram_123.svg"); - let old_mmd = media_dir.join("old_diagram_123.mmd"); let current_svg = media_dir.join("current_diagram_456.svg"); let other_file = media_dir.join("readme.txt"); fs::write(&old_svg, "old svg").expect("Should write old svg"); - fs::write(&old_mmd, "old mmd").expect("Should write old mmd"); fs::write(¤t_svg, "current svg").expect("Should write current svg"); fs::write(&other_file, "readme").expect("Should write other file"); @@ -237,7 +271,7 @@ fn test_cleanup_file_detection() { for entry in entries.flatten() { let path = entry.path(); if let Some(ext) = path.extension() { - if ext == "mmd" || ext == "svg" { + if ext == "svg" { let filename = path.file_name().unwrap().to_string_lossy(); if !referenced.contains(&filename.as_ref()) { to_cleanup.push(filename.to_string()); @@ -247,7 +281,6 @@ fn test_cleanup_file_detection() { } assert!(to_cleanup.contains(&"old_diagram_123.svg".to_string())); - assert!(to_cleanup.contains(&"old_diagram_123.mmd".to_string())); assert!(!to_cleanup.contains(&"current_diagram_456.svg".to_string())); assert!(!to_cleanup.contains(&"readme.txt".to_string())); }