aether-mcp-servers 0.8.9

//! LSP-powered rename refactoring tool

use std::cmp::Reverse;
use std::path::Path;

use lsp_types::{DocumentChangeOperation, DocumentChanges, OneOf, ResourceOp, WorkspaceEdit};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use utils::display_meta::{ToolDisplayMeta, basename};

use crate::Displayed;
use crate::lsp::common::uri_to_path;
use crate::lsp::error::LspError;
use crate::lsp::registry::LspRegistry;
use crate::lsp::render::{SourceLine, render_locations};

/// Input for the `lsp_rename` tool
#[derive(Debug, Clone, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct LspRenameInput {
    /// The file containing the symbol to rename, absolute or relative to the workspace root
    #[serde(alias = "file_path")]
    pub file_path: String,
    /// The symbol name to rename (used for position resolution if line not provided)
    pub symbol: String,
    /// The new name for the symbol
    #[serde(alias = "new_name")]
    pub new_name: String,
    /// Optional 1-indexed line hint. Stale hints fall back to automatic symbol resolution.
    #[serde(default)]
    pub line: Option<u32>,
}

/// A raw LSP text edit grouped by file path.
#[derive(Debug, Clone)]
struct LspFileEdit {
    file_path: String,
    edits: Vec<lsp_types::TextEdit>,
}

/// Output from the `lsp_rename` tool
#[derive(Debug, Clone, Default, Serialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct LspRenameOutput {
    /// The original symbol name
    pub old_name: String,
    /// The new symbol name
    pub new_name: String,
    /// Total number of edits across all files
    pub total_edits: usize,
    /// Number of files affected
    pub files_affected: usize,
    /// Edited lines grouped by file (relative to the workspace root): `path: 3, 18`
    pub changes: String,
    /// Whether the rename was successful
    pub success: bool,
    /// Error message if rename failed
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

impl LspRenameOutput {
    async fn applied(input: &LspRenameInput, changes: &[LspFileEdit], project_root: &Path) -> Displayed<Self> {
        let edited: Vec<SourceLine> = changes
            .iter()
            .flat_map(|file_edit| {
                file_edit.edits.iter().map(|edit| SourceLine::from_range(file_edit.file_path.clone(), &edit.range))
            })
            .collect();
        let total_edits = edited.len();
        let files_affected = changes.len();

        let output = Self {
            old_name: input.symbol.clone(),
            new_name: input.new_name.clone(),
            total_edits,
            files_affected,
            changes: render_locations(&edited, project_root, None).await,
            success: true,
            ..Default::default()
        };
        Displayed::new(output, rename_success_meta(input, total_edits, files_affected))
    }

    fn failed(input: &LspRenameInput, error: String) -> Displayed<Self> {
        let output = Self {
            old_name: input.symbol.clone(),
            new_name: input.new_name.clone(),
            success: false,
            error: Some(error),
            ..Default::default()
        };
        Displayed::new(output, rename_failure_meta(input))
    }
}

/// Execute the rename operation
pub async fn execute_lsp_rename(
    mut input: LspRenameInput,
    registry: &LspRegistry,
) -> Result<Displayed<LspRenameOutput>, LspError> {
    input.file_path = registry.resolve_file(&input.file_path)?;
    let resolved = registry.resolve_symbol(&input.file_path, &input.symbol, input.line).await?;

    let workspace_edit =
        resolved.client.rename(resolved.uri, resolved.line, resolved.column, input.new_name.clone()).await?;

    let Some(edit) = workspace_edit else {
        return Ok(LspRenameOutput::failed(&input, "No changes returned from LSP server".to_string()));
    };

    let changes = match collect_workspace_text_edits(&edit) {
        Ok(changes) => changes,
        Err(error) => return Ok(LspRenameOutput::failed(&input, error.to_string())),
    };
    apply_workspace_text_edits(&changes).await?;
    Ok(LspRenameOutput::applied(&input, &changes, registry.root_path()).await)
}

/// Convert LSP `WorkspaceEdit` to grouped raw text edits.
fn collect_workspace_text_edits(edit: &WorkspaceEdit) -> Result<Vec<LspFileEdit>, LspError> {
    let mut result = if let Some(doc_changes) = &edit.document_changes {
        collect_document_changes(doc_changes)?
    } else if let Some(changes) = &edit.changes {
        changes
            .iter()
            .filter(|&(_uri, text_edits)| !text_edits.is_empty())
            .map(|(uri, text_edits)| LspFileEdit { file_path: uri_to_path(uri), edits: text_edits.clone() })
            .collect()
    } else {
        Vec::new()
    };

    result.sort_by(|a, b| a.file_path.cmp(&b.file_path));
    Ok(result)
}

fn collect_document_changes(doc_changes: &DocumentChanges) -> Result<Vec<LspFileEdit>, LspError> {
    match doc_changes {
        DocumentChanges::Edits(edits) => Ok(edits
            .iter()
            .filter_map(|doc_edit| collect_text_document_edit(&doc_edit.text_document.uri, &doc_edit.edits))
            .collect()),
        DocumentChanges::Operations(ops) => {
            ops.iter().map(collect_document_change_operation).filter_map(Result::transpose).collect()
        }
    }
}

fn collect_document_change_operation(op: &DocumentChangeOperation) -> Result<Option<LspFileEdit>, LspError> {
    match op {
        DocumentChangeOperation::Edit(doc_edit) => {
            Ok(collect_text_document_edit(&doc_edit.text_document.uri, &doc_edit.edits))
        }
        DocumentChangeOperation::Op(op) => {
            let kind = match op {
                ResourceOp::Create(_) => "create",
                ResourceOp::Rename(_) => "rename",
                ResourceOp::Delete(_) => "delete",
            };
            Err(LspError::InvalidEdit(format!("Rename returned unsupported workspace operation: {kind}")))
        }
    }
}

fn collect_text_document_edit(
    uri: &lsp_types::Uri,
    edits: &[OneOf<lsp_types::TextEdit, lsp_types::AnnotatedTextEdit>],
) -> Option<LspFileEdit> {
    let edits = edits.iter().map(extract_one_of_text_edit).collect::<Vec<_>>();
    (!edits.is_empty()).then(|| LspFileEdit { file_path: uri_to_path(uri), edits })
}

fn extract_one_of_text_edit(edit: &OneOf<lsp_types::TextEdit, lsp_types::AnnotatedTextEdit>) -> lsp_types::TextEdit {
    match edit {
        OneOf::Left(edit) => edit.clone(),
        OneOf::Right(edit) => edit.text_edit.clone(),
    }
}

async fn apply_workspace_text_edits(changes: &[LspFileEdit]) -> Result<(), LspError> {
    for file_edit in changes {
        let content = tokio::fs::read_to_string(&file_edit.file_path).await?;
        let updated = apply_file_text_edits(&content, &file_edit.edits)?;
        tokio::fs::write(&file_edit.file_path, updated).await?;
    }
    Ok(())
}

fn apply_file_text_edits(content: &str, edits: &[lsp_types::TextEdit]) -> Result<String, LspError> {
    let mut edits = edits.to_vec();
    edits.sort_by_key(|e| {
        Reverse((e.range.start.line, e.range.start.character, e.range.end.line, e.range.end.character))
    });

    let mut result = content.to_string();
    for edit in edits {
        let start = lsp_position_to_byte_offset(&result, edit.range.start)?;
        let end = lsp_position_to_byte_offset(&result, edit.range.end)?;
        if start > end || end > result.len() {
            return Err(LspError::InvalidEdit("Invalid edit range produced by LSP rename".to_string()));
        }
        result.replace_range(start..end, &edit.new_text);
    }

    Ok(result)
}

fn lsp_position_to_byte_offset(content: &str, position: lsp_types::Position) -> Result<usize, LspError> {
    let target_line = usize::try_from(position.line)
        .map_err(|_| LspError::InvalidPosition(format!("Line {} out of range", position.line)))?;
    let target_character = usize::try_from(position.character)
        .map_err(|_| LspError::InvalidPosition(format!("Character {} out of range", position.character)))?;

    let mut line_start = 0usize;
    let mut current_line = 0usize;

    while current_line < target_line {
        let Some(relative_newline) = content[line_start..].find('\n') else {
            return Err(LspError::InvalidPosition(format!(
                "Line {} not found while applying rename",
                position.line + 1
            )));
        };
        line_start += relative_newline + 1;
        current_line += 1;
    }

    let line_end = content[line_start..].find('\n').map_or(content.len(), |idx| line_start + idx);
    let line = &content[line_start..line_end];

    let mut utf16_units = 0usize;
    for (byte_idx, ch) in line.char_indices() {
        if utf16_units == target_character {
            return Ok(line_start + byte_idx);
        }
        utf16_units += ch.len_utf16();
        if utf16_units > target_character {
            return Err(LspError::InvalidPosition(format!(
                "Character {} splits a UTF-16 code unit sequence on line {}",
                position.character,
                position.line + 1
            )));
        }
    }

    if utf16_units == target_character {
        Ok(line_end)
    } else {
        Err(LspError::InvalidPosition(format!(
            "Character {} out of bounds on line {}",
            position.character,
            position.line + 1
        )))
    }
}

fn rename_success_meta(input: &LspRenameInput, total_edits: usize, files_affected: usize) -> ToolDisplayMeta {
    ToolDisplayMeta::new(
        "LSP rename",
        format!("{} → {} ({} edits in {} files)", input.symbol, input.new_name, total_edits, files_affected),
    )
}

fn rename_failure_meta(input: &LspRenameInput) -> ToolDisplayMeta {
    let file = basename(&input.file_path);
    ToolDisplayMeta::new("LSP rename failed", format!("{} in {}", input.symbol, file))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lsp::test_fixtures::{text_edit, uri};
    use lsp_types::{CreateFile, OptionalVersionedTextDocumentIdentifier, Position, TextDocumentEdit};
    use std::collections::HashMap;

    fn document_edit(path: &str, edit: lsp_types::TextEdit) -> TextDocumentEdit {
        TextDocumentEdit {
            text_document: OptionalVersionedTextDocumentIdentifier { uri: uri(path), version: Some(1) },
            edits: vec![OneOf::Left(edit)],
        }
    }

    #[test]
    #[allow(clippy::mutable_key_type)]
    fn test_convert_simple_changes() {
        let changes = HashMap::from([(
            uri("/src/main.rs"),
            vec![text_edit(0, 5, 10, "new_name"), text_edit(5, 0, 8, "new_name")],
        )]);
        let edit = WorkspaceEdit { changes: Some(changes), ..Default::default() };

        let result = collect_workspace_text_edits(&edit).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].file_path, "/src/main.rs");
        assert_eq!(result[0].edits, [text_edit(0, 5, 10, "new_name"), text_edit(5, 0, 8, "new_name")]);
    }

    #[test]
    fn test_convert_document_changes() {
        let doc_edit = document_edit("/src/lib.rs", text_edit(10, 2, 6, "renamed"));
        let edit =
            WorkspaceEdit { document_changes: Some(DocumentChanges::Edits(vec![doc_edit])), ..Default::default() };

        let result = collect_workspace_text_edits(&edit).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].file_path, "/src/lib.rs");
        assert_eq!(result[0].edits, [text_edit(10, 2, 6, "renamed")]);
    }

    #[test]
    #[allow(clippy::mutable_key_type)]
    fn test_document_changes_take_precedence_over_changes() {
        let edit = WorkspaceEdit {
            changes: Some(HashMap::from([(uri("/src/changes.rs"), vec![text_edit(0, 0, 3, "old")])])),
            document_changes: Some(DocumentChanges::Edits(vec![document_edit(
                "/src/document_changes.rs",
                text_edit(1, 0, 3, "new"),
            )])),
            change_annotations: None,
        };

        let result = collect_workspace_text_edits(&edit).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].file_path, "/src/document_changes.rs");
    }

    #[test]
    fn test_rejects_unsupported_resource_operations() {
        let create = CreateFile { uri: uri("/src/new.rs"), options: None, annotation_id: None };
        let edit = WorkspaceEdit {
            document_changes: Some(DocumentChanges::Operations(vec![DocumentChangeOperation::Op(ResourceOp::Create(
                create,
            ))])),
            ..Default::default()
        };

        let err = collect_workspace_text_edits(&edit).unwrap_err().to_string();
        assert!(err.contains("unsupported workspace operation: create"), "{err}");
    }

    #[test]
    fn test_apply_file_text_edits_replaces_multiple_ranges_from_end() {
        let content = "fn greet() {\n    greet();\n    greet();\n}\n";
        let edits = [text_edit(0, 3, 8, "hello"), text_edit(1, 4, 9, "hello"), text_edit(2, 4, 9, "hello")];

        let updated = apply_file_text_edits(content, &edits).unwrap();
        assert_eq!(updated, "fn hello() {\n    hello();\n    hello();\n}\n");
    }

    #[test]
    fn test_lsp_position_to_byte_offset_handles_utf16_columns() {
        let offsets: Vec<_> = [0, 1, 3, 4]
            .into_iter()
            .map(|character| lsp_position_to_byte_offset("a😀z\n", Position::new(0, character)).unwrap())
            .collect();

        assert_eq!(offsets, [0, 1, 5, 6]);
    }
}