Fix MCP settings migration continually adding the same key (#32848)

Release Notes:

- N/A

crates/migrator/src/migrations/m_2025_06_16/settings.rs

@@ -4,149 +4,87 @@ use tree_sitter::{Query, QueryMatch};
 
 use crate::MigrationPatterns;
 
-pub const SETTINGS_PATTERNS: MigrationPatterns = &[
+pub const SETTINGS_PATTERNS: MigrationPatterns = &[(
-    (
+    SETTINGS_CONTEXT_SERVER_PATTERN,
-        SETTINGS_CUSTOM_CONTEXT_SERVER_PATTERN,
+    migrate_context_server_settings,
-        migrate_custom_context_server_settings,
+)];
-    ),
-    (
-        SETTINGS_EXTENSION_CONTEXT_SERVER_PATTERN,
-        migrate_extension_context_server_settings,
-    ),
-    (
-        SETTINGS_EMPTY_CONTEXT_SERVER_PATTERN,
-        migrate_empty_context_server_settings,
-    ),
-];
 
-const SETTINGS_CUSTOM_CONTEXT_SERVER_PATTERN: &str = r#"(document
+const SETTINGS_CONTEXT_SERVER_PATTERN: &str = r#"(document
     (object
         (pair
             key: (string (string_content) @context-servers)
             value: (object
                 (pair
-                    key: (string)
+                    key: (string (string_content) @server-name)
-                    value: (object
-                        (pair
-                            key: (string (string_content) @previous-key)
-                            value: (object)
-                        )*
-                        (pair
-                            key: (string (string_content) @key)
-                            value: (object)
-                        )
-                        (pair
-                            key: (string (string_content) @next-key)
-                            value: (object)
-                        )*
-                    ) @server-settings
-                )
-            )
-        )
-    )
-    (#eq? @context-servers "context_servers")
-    (#eq? @key "command")
-    (#not-eq? @previous-key "source")
-    (#not-eq? @next-key "source")
-)"#;
-
-fn migrate_custom_context_server_settings(
-    _contents: &str,
-    mat: &QueryMatch,
-    query: &Query,
-) -> Option<(Range<usize>, String)> {
-    let server_settings_index = query.capture_index_for_name("server-settings")?;
-    let server_settings = mat.nodes_for_capture_index(server_settings_index).next()?;
-    // Move forward 1 to get inside the object
-    let start = server_settings.start_byte() + 1;
-
-    Some((
-        start..start,
-        r#"
-            "source": "custom","#
-            .to_string(),
-    ))
-}
-
-const SETTINGS_EXTENSION_CONTEXT_SERVER_PATTERN: &str = r#"(document
-    (object
-        (pair
-            key: (string (string_content) @context-servers)
-            value: (object
-                (pair
-                    key: (string)
-                    value: (object
-                        (pair
-                            key: (string (string_content) @previous-key)
-                            value: (object)
-                        )*
-                        (pair
-                            key: (string (string_content) @key)
-                            value: (object)
-                        )
-                        (pair
-                            key: (string (string_content) @next-key)
-                            value: (object)
-                        )*
-                    ) @server-settings
-                )
-            )
-        )
-    )
-    (#eq? @context-servers "context_servers")
-    (#eq? @key "settings")
-    (#not-match? @previous-key "^command|source$")
-    (#not-match? @next-key "^command|source$")
-)"#;
-
-fn migrate_extension_context_server_settings(
-    _contents: &str,
-    mat: &QueryMatch,
-    query: &Query,
-) -> Option<(Range<usize>, String)> {
-    let server_settings_index = query.capture_index_for_name("server-settings")?;
-    let server_settings = mat.nodes_for_capture_index(server_settings_index).next()?;
-    // Move forward 1 to get inside the object
-    let start = server_settings.start_byte() + 1;
-
-    Some((
-        start..start,
-        r#"
-            "source": "extension","#
-            .to_string(),
-    ))
-}
-
-const SETTINGS_EMPTY_CONTEXT_SERVER_PATTERN: &str = r#"(document
-    (object
-        (pair
-            key: (string (string_content) @context-servers)
-            value: (object
-                (pair
-                    key: (string)
                     value: (object) @server-settings
                 )
             )
         )
     )
     (#eq? @context-servers "context_servers")
-    (#eq? @server-settings "{}")
 )"#;
 
-fn migrate_empty_context_server_settings(
+fn migrate_context_server_settings(
-    _contents: &str,
+    contents: &str,
     mat: &QueryMatch,
     query: &Query,
 ) -> Option<(Range<usize>, String)> {
     let server_settings_index = query.capture_index_for_name("server-settings")?;
     let server_settings = mat.nodes_for_capture_index(server_settings_index).next()?;
 
+    let mut has_command = false;
+    let mut has_settings = false;
+    let mut other_keys = 0;
+    let mut column = None;
+
+    // Parse the server settings to check what keys it contains
+    let mut cursor = server_settings.walk();
+    for child in server_settings.children(&mut cursor) {
+        if child.kind() == "pair" {
+            if let Some(key_node) = child.child_by_field_name("key") {
+                if let (None, Some(quote_content)) = (column, key_node.child(0)) {
+                    column = Some(quote_content.start_position().column);
+                }
+                if let Some(string_content) = key_node.child(1) {
+                    let key = &contents[string_content.byte_range()];
+                    match key {
+                        // If it already has a source key, don't modify it
+                        "source" => return None,
+                        "command" => has_command = true,
+                        "settings" => has_settings = true,
+                        _ => other_keys += 1,
+                    }
+                }
+            }
+        }
+    }
+
+    let source_type = if has_command { "custom" } else { "extension" };
+
+    // Insert the source key at the beginning of the object
+    let start = server_settings.start_byte() + 1;
+    let indent = " ".repeat(column.unwrap_or(12));
+
+    if !has_command && !has_settings {
+        return Some((
+            start..start,
+            format!(
+                r#"
+{indent}"source": "{}",
+{indent}"settings": {{}}{}
+        "#,
+                source_type,
+                if other_keys > 0 { "," } else { "" }
+            ),
+        ));
+    }
+
     Some((
-        server_settings.byte_range(),
+        start..start,
-        r#"{
+        format!(
-            "source": "extension",
+            r#"
-            "settings": {}
+{indent}"source": "{}","#,
-        }"#
+            source_type
-        .to_string(),
+        ),
     ))
 }

crates/migrator/src/migrator.rs

@@ -1010,4 +1010,46 @@ mod tests {
             ),
         );
     }
+
+    #[test]
+    fn test_mcp_settings_migration_doesnt_change_valid_settings() {
+        let settings = r#"{
+    "context_servers": {
+        "empty_server": {
+            "source": "extension",
+            "settings": {}
+        },
+        "extension_server": {
+            "source": "extension",
+            "settings": {
+                "foo": "bar"
+            }
+        },
+        "custom_server": {
+            "source": "custom",
+            "command": {
+                "path": "foo",
+                "args": ["bar"],
+                "env": {
+                    "FOO": "BAR"
+                }
+            }
+        },
+        "invalid_server": {
+            "source": "custom",
+            "command": {
+                "path": "foo",
+                "args": ["bar"],
+                "env": {
+                    "FOO": "BAR"
+                }
+            },
+            "settings": {
+                "foo": "bar"
+            }
+        }
+    }
+}"#;
+        assert_migrate_settings(settings, None);
+    }
 }