migrator.rs

   1//! ## When to create a migration and why?
   2//! A migration is necessary when keymap actions or settings are renamed or transformed (e.g., from an array to a string, a string to an array, a boolean to an enum, etc.).
   3//!
   4//! This ensures that users with outdated settings are automatically updated to use the corresponding new settings internally.
   5//! It also provides a quick way to migrate their existing settings to the latest state using button in UI.
   6//!
   7//! ## How to create a migration?
   8//! Migrations use Tree-sitter to query commonly used patterns, such as actions with a string or actions with an array where the second argument is an object, etc.
   9//! Once queried, *you can filter out the modified items* and write the replacement logic.
  10//!
  11//! You *must not* modify previous migrations; always create new ones instead.
  12//! This is important because if a user is in an intermediate state, they can smoothly transition to the latest state.
  13//! Modifying existing migrations means they will only work for users upgrading from version x-1 to x, but not from x-2 to x, and so on, where x is the latest version.
  14//!
  15//! You only need to write replacement logic for x-1 to x because you can be certain that, internally, every user will be at x-1, regardless of their on disk state.
  16
  17use anyhow::{Context as _, Result};
  18use std::{cmp::Reverse, ops::Range, sync::LazyLock};
  19use streaming_iterator::StreamingIterator;
  20use tree_sitter::{Query, QueryMatch};
  21
  22use patterns::SETTINGS_NESTED_KEY_VALUE_PATTERN;
  23
  24mod migrations;
  25mod patterns;
  26
  27fn migrate(text: &str, patterns: MigrationPatterns, query: &Query) -> Result<Option<String>> {
  28    let mut parser = tree_sitter::Parser::new();
  29    parser.set_language(&tree_sitter_json::LANGUAGE.into())?;
  30    let syntax_tree = parser
  31        .parse(&text, None)
  32        .context("failed to parse settings")?;
  33
  34    let mut cursor = tree_sitter::QueryCursor::new();
  35    let mut matches = cursor.matches(query, syntax_tree.root_node(), text.as_bytes());
  36
  37    let mut edits = vec![];
  38    while let Some(mat) = matches.next() {
  39        if let Some((_, callback)) = patterns.get(mat.pattern_index) {
  40            edits.extend(callback(&text, &mat, query));
  41        }
  42    }
  43
  44    edits.sort_by_key(|(range, _)| (range.start, Reverse(range.end)));
  45    edits.dedup_by(|(range_b, _), (range_a, _)| {
  46        range_a.contains(&range_b.start) || range_a.contains(&range_b.end)
  47    });
  48
  49    if edits.is_empty() {
  50        Ok(None)
  51    } else {
  52        let mut new_text = text.to_string();
  53        for (range, replacement) in edits.iter().rev() {
  54            new_text.replace_range(range.clone(), replacement);
  55        }
  56        if new_text == text {
  57            log::error!(
  58                "Edits computed for configuration migration do not cause a change: {:?}",
  59                edits
  60            );
  61            Ok(None)
  62        } else {
  63            Ok(Some(new_text))
  64        }
  65    }
  66}
  67
  68fn run_migrations(
  69    text: &str,
  70    migrations: &[(MigrationPatterns, &Query)],
  71) -> Result<Option<String>> {
  72    let mut current_text = text.to_string();
  73    let mut result: Option<String> = None;
  74    for (patterns, query) in migrations.iter() {
  75        if let Some(migrated_text) = migrate(&current_text, patterns, query)? {
  76            current_text = migrated_text.clone();
  77            result = Some(migrated_text);
  78        }
  79    }
  80    Ok(result.filter(|new_text| text != new_text))
  81}
  82
  83pub fn migrate_keymap(text: &str) -> Result<Option<String>> {
  84    let migrations: &[(MigrationPatterns, &Query)] = &[
  85        (
  86            migrations::m_2025_01_29::KEYMAP_PATTERNS,
  87            &KEYMAP_QUERY_2025_01_29,
  88        ),
  89        (
  90            migrations::m_2025_01_30::KEYMAP_PATTERNS,
  91            &KEYMAP_QUERY_2025_01_30,
  92        ),
  93        (
  94            migrations::m_2025_03_03::KEYMAP_PATTERNS,
  95            &KEYMAP_QUERY_2025_03_03,
  96        ),
  97        (
  98            migrations::m_2025_03_06::KEYMAP_PATTERNS,
  99            &KEYMAP_QUERY_2025_03_06,
 100        ),
 101        (
 102            migrations::m_2025_04_15::KEYMAP_PATTERNS,
 103            &KEYMAP_QUERY_2025_04_15,
 104        ),
 105    ];
 106    run_migrations(text, migrations)
 107}
 108
 109pub fn migrate_settings(text: &str) -> Result<Option<String>> {
 110    let migrations: &[(MigrationPatterns, &Query)] = &[
 111        (
 112            migrations::m_2025_01_02::SETTINGS_PATTERNS,
 113            &SETTINGS_QUERY_2025_01_02,
 114        ),
 115        (
 116            migrations::m_2025_01_29::SETTINGS_PATTERNS,
 117            &SETTINGS_QUERY_2025_01_29,
 118        ),
 119        (
 120            migrations::m_2025_01_30::SETTINGS_PATTERNS,
 121            &SETTINGS_QUERY_2025_01_30,
 122        ),
 123        (
 124            migrations::m_2025_03_29::SETTINGS_PATTERNS,
 125            &SETTINGS_QUERY_2025_03_29,
 126        ),
 127        (
 128            migrations::m_2025_04_15::SETTINGS_PATTERNS,
 129            &SETTINGS_QUERY_2025_04_15,
 130        ),
 131        (
 132            migrations::m_2025_04_21::SETTINGS_PATTERNS,
 133            &SETTINGS_QUERY_2025_04_21,
 134        ),
 135        (
 136            migrations::m_2025_04_23::SETTINGS_PATTERNS,
 137            &SETTINGS_QUERY_2025_04_23,
 138        ),
 139        (
 140            migrations::m_2025_05_05::SETTINGS_PATTERNS,
 141            &SETTINGS_QUERY_2025_05_05,
 142        ),
 143        (
 144            migrations::m_2025_05_08::SETTINGS_PATTERNS,
 145            &SETTINGS_QUERY_2025_05_08,
 146        ),
 147        (
 148            migrations::m_2025_05_29::SETTINGS_PATTERNS,
 149            &SETTINGS_QUERY_2025_05_29,
 150        ),
 151        (
 152            migrations::m_2025_06_16::SETTINGS_PATTERNS,
 153            &SETTINGS_QUERY_2025_06_16,
 154        ),
 155        (
 156            migrations::m_2025_06_27::SETTINGS_PATTERNS,
 157            &SETTINGS_QUERY_2025_06_27,
 158        ),
 159    ];
 160    run_migrations(text, migrations)
 161}
 162
 163pub fn migrate_edit_prediction_provider_settings(text: &str) -> Result<Option<String>> {
 164    migrate(
 165        &text,
 166        &[(
 167            SETTINGS_NESTED_KEY_VALUE_PATTERN,
 168            migrations::m_2025_01_29::replace_edit_prediction_provider_setting,
 169        )],
 170        &EDIT_PREDICTION_SETTINGS_MIGRATION_QUERY,
 171    )
 172}
 173
 174pub type MigrationPatterns = &'static [(
 175    &'static str,
 176    fn(&str, &QueryMatch, &Query) -> Option<(Range<usize>, String)>,
 177)];
 178
 179macro_rules! define_query {
 180    ($var_name:ident, $patterns_path:path) => {
 181        static $var_name: LazyLock<Query> = LazyLock::new(|| {
 182            Query::new(
 183                &tree_sitter_json::LANGUAGE.into(),
 184                &$patterns_path
 185                    .iter()
 186                    .map(|pattern| pattern.0)
 187                    .collect::<String>(),
 188            )
 189            .unwrap()
 190        });
 191    };
 192}
 193
 194// keymap
 195define_query!(
 196    KEYMAP_QUERY_2025_01_29,
 197    migrations::m_2025_01_29::KEYMAP_PATTERNS
 198);
 199define_query!(
 200    KEYMAP_QUERY_2025_01_30,
 201    migrations::m_2025_01_30::KEYMAP_PATTERNS
 202);
 203define_query!(
 204    KEYMAP_QUERY_2025_03_03,
 205    migrations::m_2025_03_03::KEYMAP_PATTERNS
 206);
 207define_query!(
 208    KEYMAP_QUERY_2025_03_06,
 209    migrations::m_2025_03_06::KEYMAP_PATTERNS
 210);
 211define_query!(
 212    KEYMAP_QUERY_2025_04_15,
 213    migrations::m_2025_04_15::KEYMAP_PATTERNS
 214);
 215
 216// settings
 217define_query!(
 218    SETTINGS_QUERY_2025_01_02,
 219    migrations::m_2025_01_02::SETTINGS_PATTERNS
 220);
 221define_query!(
 222    SETTINGS_QUERY_2025_01_29,
 223    migrations::m_2025_01_29::SETTINGS_PATTERNS
 224);
 225define_query!(
 226    SETTINGS_QUERY_2025_01_30,
 227    migrations::m_2025_01_30::SETTINGS_PATTERNS
 228);
 229define_query!(
 230    SETTINGS_QUERY_2025_03_29,
 231    migrations::m_2025_03_29::SETTINGS_PATTERNS
 232);
 233define_query!(
 234    SETTINGS_QUERY_2025_04_15,
 235    migrations::m_2025_04_15::SETTINGS_PATTERNS
 236);
 237define_query!(
 238    SETTINGS_QUERY_2025_04_21,
 239    migrations::m_2025_04_21::SETTINGS_PATTERNS
 240);
 241define_query!(
 242    SETTINGS_QUERY_2025_04_23,
 243    migrations::m_2025_04_23::SETTINGS_PATTERNS
 244);
 245define_query!(
 246    SETTINGS_QUERY_2025_05_05,
 247    migrations::m_2025_05_05::SETTINGS_PATTERNS
 248);
 249define_query!(
 250    SETTINGS_QUERY_2025_05_08,
 251    migrations::m_2025_05_08::SETTINGS_PATTERNS
 252);
 253define_query!(
 254    SETTINGS_QUERY_2025_05_29,
 255    migrations::m_2025_05_29::SETTINGS_PATTERNS
 256);
 257define_query!(
 258    SETTINGS_QUERY_2025_06_16,
 259    migrations::m_2025_06_16::SETTINGS_PATTERNS
 260);
 261define_query!(
 262    SETTINGS_QUERY_2025_06_27,
 263    migrations::m_2025_06_27::SETTINGS_PATTERNS
 264);
 265
 266// custom query
 267static EDIT_PREDICTION_SETTINGS_MIGRATION_QUERY: LazyLock<Query> = LazyLock::new(|| {
 268    Query::new(
 269        &tree_sitter_json::LANGUAGE.into(),
 270        SETTINGS_NESTED_KEY_VALUE_PATTERN,
 271    )
 272    .unwrap()
 273});
 274
 275#[cfg(test)]
 276mod tests {
 277    use super::*;
 278
 279    fn assert_migrate_keymap(input: &str, output: Option<&str>) {
 280        let migrated = migrate_keymap(&input).unwrap();
 281        pretty_assertions::assert_eq!(migrated.as_deref(), output);
 282    }
 283
 284    fn assert_migrate_settings(input: &str, output: Option<&str>) {
 285        let migrated = migrate_settings(&input).unwrap();
 286        pretty_assertions::assert_eq!(migrated.as_deref(), output);
 287    }
 288
 289    fn assert_migrate_settings_with_migrations(
 290        migrations: &[(MigrationPatterns, &Query)],
 291        input: &str,
 292        output: Option<&str>,
 293    ) {
 294        let migrated = run_migrations(input, migrations).unwrap();
 295        pretty_assertions::assert_eq!(migrated.as_deref(), output);
 296    }
 297
 298    #[test]
 299    fn test_replace_array_with_single_string() {
 300        assert_migrate_keymap(
 301            r#"
 302            [
 303                {
 304                    "bindings": {
 305                        "cmd-1": ["workspace::ActivatePaneInDirection", "Up"]
 306                    }
 307                }
 308            ]
 309            "#,
 310            Some(
 311                r#"
 312            [
 313                {
 314                    "bindings": {
 315                        "cmd-1": "workspace::ActivatePaneUp"
 316                    }
 317                }
 318            ]
 319            "#,
 320            ),
 321        )
 322    }
 323
 324    #[test]
 325    fn test_replace_action_argument_object_with_single_value() {
 326        assert_migrate_keymap(
 327            r#"
 328            [
 329                {
 330                    "bindings": {
 331                        "cmd-1": ["editor::FoldAtLevel", { "level": 1 }]
 332                    }
 333                }
 334            ]
 335            "#,
 336            Some(
 337                r#"
 338            [
 339                {
 340                    "bindings": {
 341                        "cmd-1": ["editor::FoldAtLevel", 1]
 342                    }
 343                }
 344            ]
 345            "#,
 346            ),
 347        )
 348    }
 349
 350    #[test]
 351    fn test_replace_action_argument_object_with_single_value_2() {
 352        assert_migrate_keymap(
 353            r#"
 354            [
 355                {
 356                    "bindings": {
 357                        "cmd-1": ["vim::PushOperator", { "Object": { "some" : "value" } }]
 358                    }
 359                }
 360            ]
 361            "#,
 362            Some(
 363                r#"
 364            [
 365                {
 366                    "bindings": {
 367                        "cmd-1": ["vim::PushObject", { "some" : "value" }]
 368                    }
 369                }
 370            ]
 371            "#,
 372            ),
 373        )
 374    }
 375
 376    #[test]
 377    fn test_rename_string_action() {
 378        assert_migrate_keymap(
 379            r#"
 380                [
 381                    {
 382                        "bindings": {
 383                            "cmd-1": "inline_completion::ToggleMenu"
 384                        }
 385                    }
 386                ]
 387            "#,
 388            Some(
 389                r#"
 390                [
 391                    {
 392                        "bindings": {
 393                            "cmd-1": "edit_prediction::ToggleMenu"
 394                        }
 395                    }
 396                ]
 397            "#,
 398            ),
 399        )
 400    }
 401
 402    #[test]
 403    fn test_rename_context_key() {
 404        assert_migrate_keymap(
 405            r#"
 406                [
 407                    {
 408                        "context": "Editor && inline_completion && !showing_completions"
 409                    }
 410                ]
 411            "#,
 412            Some(
 413                r#"
 414                [
 415                    {
 416                        "context": "Editor && edit_prediction && !showing_completions"
 417                    }
 418                ]
 419            "#,
 420            ),
 421        )
 422    }
 423
 424    #[test]
 425    fn test_incremental_migrations() {
 426        // Here string transforms to array internally. Then, that array transforms back to string.
 427        assert_migrate_keymap(
 428            r#"
 429                [
 430                    {
 431                        "bindings": {
 432                            "ctrl-q": "editor::GoToHunk", // should remain same
 433                            "ctrl-w": "editor::GoToPrevHunk", // should rename
 434                            "ctrl-q": ["editor::GoToHunk", { "center_cursor": true }], // should transform
 435                            "ctrl-w": ["editor::GoToPreviousHunk", { "center_cursor": true }] // should transform
 436                        }
 437                    }
 438                ]
 439            "#,
 440            Some(
 441                r#"
 442                [
 443                    {
 444                        "bindings": {
 445                            "ctrl-q": "editor::GoToHunk", // should remain same
 446                            "ctrl-w": "editor::GoToPreviousHunk", // should rename
 447                            "ctrl-q": "editor::GoToHunk", // should transform
 448                            "ctrl-w": "editor::GoToPreviousHunk" // should transform
 449                        }
 450                    }
 451                ]
 452            "#,
 453            ),
 454        )
 455    }
 456
 457    #[test]
 458    fn test_action_argument_snake_case() {
 459        // First performs transformations, then replacements
 460        assert_migrate_keymap(
 461            r#"
 462            [
 463                {
 464                    "bindings": {
 465                        "cmd-1": ["vim::PushOperator", { "Object": { "around": false } }],
 466                        "cmd-3": ["pane::CloseActiveItem", { "saveIntent": "saveAll" }],
 467                        "cmd-2": ["vim::NextWordStart", { "ignorePunctuation": true }],
 468                        "cmd-4": ["task::Spawn", { "task_name": "a b" }] // should remain as it is
 469                    }
 470                }
 471            ]
 472            "#,
 473            Some(
 474                r#"
 475            [
 476                {
 477                    "bindings": {
 478                        "cmd-1": ["vim::PushObject", { "around": false }],
 479                        "cmd-3": ["pane::CloseActiveItem", { "save_intent": "save_all" }],
 480                        "cmd-2": ["vim::NextWordStart", { "ignore_punctuation": true }],
 481                        "cmd-4": ["task::Spawn", { "task_name": "a b" }] // should remain as it is
 482                    }
 483                }
 484            ]
 485            "#,
 486            ),
 487        )
 488    }
 489
 490    #[test]
 491    fn test_replace_setting_name() {
 492        assert_migrate_settings(
 493            r#"
 494                {
 495                    "show_inline_completions_in_menu": true,
 496                    "show_inline_completions": true,
 497                    "inline_completions_disabled_in": ["string"],
 498                    "inline_completions": { "some" : "value" }
 499                }
 500            "#,
 501            Some(
 502                r#"
 503                {
 504                    "show_edit_predictions_in_menu": true,
 505                    "show_edit_predictions": true,
 506                    "edit_predictions_disabled_in": ["string"],
 507                    "edit_predictions": { "some" : "value" }
 508                }
 509            "#,
 510            ),
 511        )
 512    }
 513
 514    #[test]
 515    fn test_nested_string_replace_for_settings() {
 516        assert_migrate_settings(
 517            r#"
 518                {
 519                    "features": {
 520                        "inline_completion_provider": "zed"
 521                    },
 522                }
 523            "#,
 524            Some(
 525                r#"
 526                {
 527                    "features": {
 528                        "edit_prediction_provider": "zed"
 529                    },
 530                }
 531            "#,
 532            ),
 533        )
 534    }
 535
 536    #[test]
 537    fn test_replace_settings_in_languages() {
 538        assert_migrate_settings(
 539            r#"
 540                {
 541                    "languages": {
 542                        "Astro": {
 543                            "show_inline_completions": true
 544                        }
 545                    }
 546                }
 547            "#,
 548            Some(
 549                r#"
 550                {
 551                    "languages": {
 552                        "Astro": {
 553                            "show_edit_predictions": true
 554                        }
 555                    }
 556                }
 557            "#,
 558            ),
 559        )
 560    }
 561
 562    #[test]
 563    fn test_replace_settings_value() {
 564        assert_migrate_settings(
 565            r#"
 566                {
 567                    "scrollbar": {
 568                        "diagnostics": true
 569                    },
 570                    "chat_panel": {
 571                        "button": true
 572                    }
 573                }
 574            "#,
 575            Some(
 576                r#"
 577                {
 578                    "scrollbar": {
 579                        "diagnostics": "all"
 580                    },
 581                    "chat_panel": {
 582                        "button": "always"
 583                    }
 584                }
 585            "#,
 586            ),
 587        )
 588    }
 589
 590    #[test]
 591    fn test_replace_settings_name_and_value() {
 592        assert_migrate_settings(
 593            r#"
 594                {
 595                    "tabs": {
 596                        "always_show_close_button": true
 597                    }
 598                }
 599            "#,
 600            Some(
 601                r#"
 602                {
 603                    "tabs": {
 604                        "show_close_button": "always"
 605                    }
 606                }
 607            "#,
 608            ),
 609        )
 610    }
 611
 612    #[test]
 613    fn test_replace_bash_with_terminal_in_profiles() {
 614        assert_migrate_settings(
 615            r#"
 616                {
 617                    "assistant": {
 618                        "profiles": {
 619                            "custom": {
 620                                "name": "Custom",
 621                                "tools": {
 622                                    "bash": true,
 623                                    "diagnostics": true
 624                                }
 625                            }
 626                        }
 627                    }
 628                }
 629            "#,
 630            Some(
 631                r#"
 632                {
 633                    "agent": {
 634                        "profiles": {
 635                            "custom": {
 636                                "name": "Custom",
 637                                "tools": {
 638                                    "terminal": true,
 639                                    "diagnostics": true
 640                                }
 641                            }
 642                        }
 643                    }
 644                }
 645            "#,
 646            ),
 647        )
 648    }
 649
 650    #[test]
 651    fn test_replace_bash_false_with_terminal_in_profiles() {
 652        assert_migrate_settings(
 653            r#"
 654                {
 655                    "assistant": {
 656                        "profiles": {
 657                            "custom": {
 658                                "name": "Custom",
 659                                "tools": {
 660                                    "bash": false,
 661                                    "diagnostics": true
 662                                }
 663                            }
 664                        }
 665                    }
 666                }
 667            "#,
 668            Some(
 669                r#"
 670                {
 671                    "agent": {
 672                        "profiles": {
 673                            "custom": {
 674                                "name": "Custom",
 675                                "tools": {
 676                                    "terminal": false,
 677                                    "diagnostics": true
 678                                }
 679                            }
 680                        }
 681                    }
 682                }
 683            "#,
 684            ),
 685        )
 686    }
 687
 688    #[test]
 689    fn test_no_bash_in_profiles() {
 690        assert_migrate_settings(
 691            r#"
 692                {
 693                    "assistant": {
 694                        "profiles": {
 695                            "custom": {
 696                                "name": "Custom",
 697                                "tools": {
 698                                    "diagnostics": true,
 699                                    "find_path": true,
 700                                    "read_file": true
 701                                }
 702                            }
 703                        }
 704                    }
 705                }
 706            "#,
 707            Some(
 708                r#"
 709                {
 710                    "agent": {
 711                        "profiles": {
 712                            "custom": {
 713                                "name": "Custom",
 714                                "tools": {
 715                                    "diagnostics": true,
 716                                    "find_path": true,
 717                                    "read_file": true
 718                                }
 719                            }
 720                        }
 721                    }
 722                }
 723            "#,
 724            ),
 725        )
 726    }
 727
 728    #[test]
 729    fn test_rename_path_search_to_find_path() {
 730        assert_migrate_settings(
 731            r#"
 732                {
 733                    "assistant": {
 734                        "profiles": {
 735                            "default": {
 736                                "tools": {
 737                                    "path_search": true,
 738                                    "read_file": true
 739                                }
 740                            }
 741                        }
 742                    }
 743                }
 744            "#,
 745            Some(
 746                r#"
 747                {
 748                    "agent": {
 749                        "profiles": {
 750                            "default": {
 751                                "tools": {
 752                                    "find_path": true,
 753                                    "read_file": true
 754                                }
 755                            }
 756                        }
 757                    }
 758                }
 759            "#,
 760            ),
 761        );
 762    }
 763
 764    #[test]
 765    fn test_rename_assistant() {
 766        assert_migrate_settings(
 767            r#"{
 768                "assistant": {
 769                    "foo": "bar"
 770                },
 771                "edit_predictions": {
 772                    "enabled_in_assistant": false,
 773                }
 774            }"#,
 775            Some(
 776                r#"{
 777                "agent": {
 778                    "foo": "bar"
 779                },
 780                "edit_predictions": {
 781                    "enabled_in_text_threads": false,
 782                }
 783            }"#,
 784            ),
 785        );
 786    }
 787
 788    #[test]
 789    fn test_comment_duplicated_agent() {
 790        assert_migrate_settings(
 791            r#"{
 792                "agent": {
 793                    "name": "assistant-1",
 794                "model": "gpt-4", // weird formatting
 795                    "utf8": "привіт"
 796                },
 797                "something": "else",
 798                "agent": {
 799                    "name": "assistant-2",
 800                    "model": "gemini-pro"
 801                }
 802            }
 803        "#,
 804            Some(
 805                r#"{
 806                /* Duplicated key auto-commented: "agent": {
 807                    "name": "assistant-1",
 808                "model": "gpt-4", // weird formatting
 809                    "utf8": "привіт"
 810                }, */
 811                "something": "else",
 812                "agent": {
 813                    "name": "assistant-2",
 814                    "model": "gemini-pro"
 815                }
 816            }
 817        "#,
 818            ),
 819        );
 820    }
 821
 822    #[test]
 823    fn test_preferred_completion_mode_migration() {
 824        assert_migrate_settings(
 825            r#"{
 826                "agent": {
 827                    "preferred_completion_mode": "max",
 828                    "enabled": true
 829                }
 830            }"#,
 831            Some(
 832                r#"{
 833                "agent": {
 834                    "preferred_completion_mode": "burn",
 835                    "enabled": true
 836                }
 837            }"#,
 838            ),
 839        );
 840
 841        assert_migrate_settings(
 842            r#"{
 843                "agent": {
 844                    "preferred_completion_mode": "normal",
 845                    "enabled": true
 846                }
 847            }"#,
 848            None,
 849        );
 850
 851        assert_migrate_settings(
 852            r#"{
 853                "agent": {
 854                    "preferred_completion_mode": "burn",
 855                    "enabled": true
 856                }
 857            }"#,
 858            None,
 859        );
 860
 861        assert_migrate_settings(
 862            r#"{
 863                "other_section": {
 864                    "preferred_completion_mode": "max"
 865                },
 866                "agent": {
 867                    "preferred_completion_mode": "max"
 868                }
 869            }"#,
 870            Some(
 871                r#"{
 872                "other_section": {
 873                    "preferred_completion_mode": "max"
 874                },
 875                "agent": {
 876                    "preferred_completion_mode": "burn"
 877                }
 878            }"#,
 879            ),
 880        );
 881    }
 882
 883    #[test]
 884    fn test_mcp_settings_migration() {
 885        assert_migrate_settings_with_migrations(
 886            &[(
 887                migrations::m_2025_06_16::SETTINGS_PATTERNS,
 888                &SETTINGS_QUERY_2025_06_16,
 889            )],
 890            r#"{
 891    "context_servers": {
 892        "empty_server": {},
 893        "extension_server": {
 894            "settings": {
 895                "foo": "bar"
 896            }
 897        },
 898        "custom_server": {
 899            "command": {
 900                "path": "foo",
 901                "args": ["bar"],
 902                "env": {
 903                    "FOO": "BAR"
 904                }
 905            }
 906        },
 907        "invalid_server": {
 908            "command": {
 909                "path": "foo",
 910                "args": ["bar"],
 911                "env": {
 912                    "FOO": "BAR"
 913                }
 914            },
 915            "settings": {
 916                "foo": "bar"
 917            }
 918        },
 919        "empty_server2": {},
 920        "extension_server2": {
 921            "foo": "bar",
 922            "settings": {
 923                "foo": "bar"
 924            },
 925            "bar": "foo"
 926        },
 927        "custom_server2": {
 928            "foo": "bar",
 929            "command": {
 930                "path": "foo",
 931                "args": ["bar"],
 932                "env": {
 933                    "FOO": "BAR"
 934                }
 935            },
 936            "bar": "foo"
 937        },
 938        "invalid_server2": {
 939            "foo": "bar",
 940            "command": {
 941                "path": "foo",
 942                "args": ["bar"],
 943                "env": {
 944                    "FOO": "BAR"
 945                }
 946            },
 947            "bar": "foo",
 948            "settings": {
 949                "foo": "bar"
 950            }
 951        }
 952    }
 953}"#,
 954            Some(
 955                r#"{
 956    "context_servers": {
 957        "empty_server": {
 958            "source": "extension",
 959            "settings": {}
 960        },
 961        "extension_server": {
 962            "source": "extension",
 963            "settings": {
 964                "foo": "bar"
 965            }
 966        },
 967        "custom_server": {
 968            "source": "custom",
 969            "command": {
 970                "path": "foo",
 971                "args": ["bar"],
 972                "env": {
 973                    "FOO": "BAR"
 974                }
 975            }
 976        },
 977        "invalid_server": {
 978            "source": "custom",
 979            "command": {
 980                "path": "foo",
 981                "args": ["bar"],
 982                "env": {
 983                    "FOO": "BAR"
 984                }
 985            },
 986            "settings": {
 987                "foo": "bar"
 988            }
 989        },
 990        "empty_server2": {
 991            "source": "extension",
 992            "settings": {}
 993        },
 994        "extension_server2": {
 995            "source": "extension",
 996            "foo": "bar",
 997            "settings": {
 998                "foo": "bar"
 999            },
1000            "bar": "foo"
1001        },
1002        "custom_server2": {
1003            "source": "custom",
1004            "foo": "bar",
1005            "command": {
1006                "path": "foo",
1007                "args": ["bar"],
1008                "env": {
1009                    "FOO": "BAR"
1010                }
1011            },
1012            "bar": "foo"
1013        },
1014        "invalid_server2": {
1015            "source": "custom",
1016            "foo": "bar",
1017            "command": {
1018                "path": "foo",
1019                "args": ["bar"],
1020                "env": {
1021                    "FOO": "BAR"
1022                }
1023            },
1024            "bar": "foo",
1025            "settings": {
1026                "foo": "bar"
1027            }
1028        }
1029    }
1030}"#,
1031            ),
1032        );
1033    }
1034
1035    #[test]
1036    fn test_mcp_settings_migration_doesnt_change_valid_settings() {
1037        let settings = r#"{
1038    "context_servers": {
1039        "empty_server": {
1040            "source": "extension",
1041            "settings": {}
1042        },
1043        "extension_server": {
1044            "source": "extension",
1045            "settings": {
1046                "foo": "bar"
1047            }
1048        },
1049        "custom_server": {
1050            "source": "custom",
1051            "command": {
1052                "path": "foo",
1053                "args": ["bar"],
1054                "env": {
1055                    "FOO": "BAR"
1056                }
1057            }
1058        },
1059        "invalid_server": {
1060            "source": "custom",
1061            "command": {
1062                "path": "foo",
1063                "args": ["bar"],
1064                "env": {
1065                    "FOO": "BAR"
1066                }
1067            },
1068            "settings": {
1069                "foo": "bar"
1070            }
1071        }
1072    }
1073}"#;
1074        assert_migrate_settings_with_migrations(
1075            &[(
1076                migrations::m_2025_06_16::SETTINGS_PATTERNS,
1077                &SETTINGS_QUERY_2025_06_16,
1078            )],
1079            settings,
1080            None,
1081        );
1082    }
1083
1084    #[test]
1085    fn test_flatten_context_server_command() {
1086        assert_migrate_settings(
1087            r#"{
1088    "context_servers": {
1089        "some-mcp-server": {
1090            "source": "custom",
1091            "command": {
1092                "path": "npx",
1093                "args": [
1094                    "-y",
1095                    "@supabase/mcp-server-supabase@latest",
1096                    "--read-only",
1097                    "--project-ref=<project-ref>"
1098                ],
1099                "env": {
1100                    "SUPABASE_ACCESS_TOKEN": "<personal-access-token>"
1101                }
1102            }
1103        }
1104    }
1105}"#,
1106            Some(
1107                r#"{
1108    "context_servers": {
1109        "some-mcp-server": {
1110            "source": "custom",
1111            "command": "npx",
1112            "args": [
1113                "-y",
1114                "@supabase/mcp-server-supabase@latest",
1115                "--read-only",
1116                "--project-ref=<project-ref>"
1117            ],
1118            "env": {
1119                "SUPABASE_ACCESS_TOKEN": "<personal-access-token>"
1120            }
1121        }
1122    }
1123}"#,
1124            ),
1125        );
1126
1127        // Test with additional keys in server object
1128        assert_migrate_settings(
1129            r#"{
1130    "context_servers": {
1131        "server-with-extras": {
1132            "source": "custom",
1133            "command": {
1134                "path": "/usr/bin/node",
1135                "args": ["server.js"]
1136            },
1137            "settings": {}
1138        }
1139    }
1140}"#,
1141            Some(
1142                r#"{
1143    "context_servers": {
1144        "server-with-extras": {
1145            "source": "custom",
1146            "command": "/usr/bin/node",
1147            "args": ["server.js"],
1148            "settings": {}
1149        }
1150    }
1151}"#,
1152            ),
1153        );
1154
1155        // Test command without args or env
1156        assert_migrate_settings(
1157            r#"{
1158    "context_servers": {
1159        "simple-server": {
1160            "source": "custom",
1161            "command": {
1162                "path": "simple-mcp-server"
1163            }
1164        }
1165    }
1166}"#,
1167            Some(
1168                r#"{
1169    "context_servers": {
1170        "simple-server": {
1171            "source": "custom",
1172            "command": "simple-mcp-server"
1173        }
1174    }
1175}"#,
1176            ),
1177        );
1178    }
1179}