agent: Rework context server settings (#32793)

This changes the way context servers are organised. We now store a
`source` which indicates if the MCP server is configured manually or
managed by an extension.

Release Notes:

- N/A

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
This commit is contained in:
Bennet Bo Fenner 2025-06-16 17:31:31 +02:00 committed by GitHub
parent c35f22dde0
commit d7db4d4e0a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 639 additions and 197 deletions

View file

@ -75,3 +75,9 @@ pub(crate) mod m_2025_05_29 {
pub(crate) use settings::SETTINGS_PATTERNS;
}
pub(crate) mod m_2025_06_16 {
mod settings;
pub(crate) use settings::SETTINGS_PATTERNS;
}

View file

@ -0,0 +1,152 @@
use std::ops::Range;
use tree_sitter::{Query, QueryMatch};
use crate::MigrationPatterns;
pub const SETTINGS_PATTERNS: MigrationPatterns = &[
(
SETTINGS_CUSTOM_CONTEXT_SERVER_PATTERN,
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
(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 "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(
_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()?;
Some((
server_settings.byte_range(),
r#"{
"source": "extension",
"settings": {}
}"#
.to_string(),
))
}

View file

@ -148,6 +148,10 @@ pub fn migrate_settings(text: &str) -> Result<Option<String>> {
migrations::m_2025_05_29::SETTINGS_PATTERNS,
&SETTINGS_QUERY_2025_05_29,
),
(
migrations::m_2025_06_16::SETTINGS_PATTERNS,
&SETTINGS_QUERY_2025_06_16,
),
];
run_migrations(text, migrations)
}
@ -246,6 +250,10 @@ define_query!(
SETTINGS_QUERY_2025_05_29,
migrations::m_2025_05_29::SETTINGS_PATTERNS
);
define_query!(
SETTINGS_QUERY_2025_06_16,
migrations::m_2025_06_16::SETTINGS_PATTERNS
);
// custom query
static EDIT_PREDICTION_SETTINGS_MIGRATION_QUERY: LazyLock<Query> = LazyLock::new(|| {
@ -854,4 +862,152 @@ mod tests {
),
);
}
#[test]
fn test_mcp_settings_migration() {
assert_migrate_settings(
r#"{
"context_servers": {
"empty_server": {},
"extension_server": {
"settings": {
"foo": "bar"
}
},
"custom_server": {
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
}
},
"invalid_server": {
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
},
"settings": {
"foo": "bar"
}
},
"empty_server2": {},
"extension_server2": {
"foo": "bar",
"settings": {
"foo": "bar"
},
"bar": "foo"
},
"custom_server2": {
"foo": "bar",
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
},
"bar": "foo"
},
"invalid_server2": {
"foo": "bar",
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
},
"bar": "foo",
"settings": {
"foo": "bar"
}
}
}
}"#,
Some(
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"
}
},
"empty_server2": {
"source": "extension",
"settings": {}
},
"extension_server2": {
"source": "extension",
"foo": "bar",
"settings": {
"foo": "bar"
},
"bar": "foo"
},
"custom_server2": {
"source": "custom",
"foo": "bar",
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
},
"bar": "foo"
},
"invalid_server2": {
"source": "custom",
"foo": "bar",
"command": {
"path": "foo",
"args": ["bar"],
"env": {
"FOO": "BAR"
}
},
"bar": "foo",
"settings": {
"foo": "bar"
}
}
}
}"#,
),
);
}
}