Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Expose context server settings to extensions (#20555)

Browse source 
This PR exposes context server settings to extensions.

Extensions can use `ContextServerSettings::for_project` to get the
context server settings for the current project.

The `experimental.context_servers` setting has been removed and replaced
with the `context_servers` setting (which is now an object instead of an
array).

Release Notes:

- N/A

---------

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Marshall Bowers 2024-11-12 17:21:58 -05:00 • committed by GitHub
parent 0a9c78a58d
commit 3ebb64ea1d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 243 additions and 126 deletions
Download patch file Download diff file Expand all files Collapse all files

2
Cargo.lock generated
View file

@ -2823,6 +2823,7 @@ dependencies = [
"log",
"parking_lot",
"postage",
"project",
"schemars",
"serde",
"serde_json",
@ -4170,6 +4171,7 @@ dependencies = [
"async-trait",
"client",
"collections",
"context_servers",
"ctor",
"env_logger 0.11.5",
"extension",

13
assets/settings/default.json
View file

@ -1182,15 +1182,6 @@
// }
// ]
"ssh_connections": [],
// Configures the Context Server Protocol binaries
//
// Examples:
// {
// "id": "server-1",
// "executable": "/path",
// "args": ['arg1", "args2"]
// }
"experimental.context_servers": {
"servers": []
}
// Configures context servers for use in the Assistant.
"context_servers": {}
}

77
crates/assistant/src/context_store.rs
View file

@ -145,49 +145,52 @@ impl ContextStore {
project: project.clone(),
prompt_builder,
};
this.handle_project_changed(project, cx);
this.handle_project_changed(project.clone(), cx);
this.synchronize_contexts(cx);
this.register_context_server_handlers(cx);
// TODO: At the time when we construct the `ContextStore` we may not have yet initialized the extensions.
// In order to register the context servers when the extension is loaded, we're periodically looping to
// see if there are context servers to register.
//
// I tried doing this in a subscription on the `ExtensionStore`, but it never seemed to fire.
//
// We should find a more elegant way to do this.
let context_server_factory_registry =
ContextServerFactoryRegistry::default_global(cx);
cx.spawn(|context_store, mut cx| async move {
loop {
let mut servers_to_register = Vec::new();
for (_id, factory) in
context_server_factory_registry.context_server_factories()
{
if let Some(server) = factory(&cx).await.log_err() {
servers_to_register.push(server);
if project.read(cx).is_local() {
// TODO: At the time when we construct the `ContextStore` we may not have yet initialized the extensions.
// In order to register the context servers when the extension is loaded, we're periodically looping to
// see if there are context servers to register.
//
// I tried doing this in a subscription on the `ExtensionStore`, but it never seemed to fire.
//
// We should find a more elegant way to do this.
let context_server_factory_registry =
ContextServerFactoryRegistry::default_global(cx);
cx.spawn(|context_store, mut cx| async move {
loop {
let mut servers_to_register = Vec::new();
for (_id, factory) in
context_server_factory_registry.context_server_factories()
{
if let Some(server) = factory(project.clone(), &cx).await.log_err()
{
servers_to_register.push(server);
}
}
let Some(_) = context_store
.update(&mut cx, |this, cx| {
this.context_server_manager.update(cx, |this, cx| {
for server in servers_to_register {
this.add_server(server, cx).detach_and_log_err(cx);
}
})
})
.log_err()
else {
break;
};
smol::Timer::after(Duration::from_millis(100)).await;
}
let Some(_) = context_store
.update(&mut cx, |this, cx| {
this.context_server_manager.update(cx, |this, cx| {
for server in servers_to_register {
this.add_server(server, cx).detach_and_log_err(cx);
}
})
})
.log_err()
else {
break;
};
smol::Timer::after(Duration::from_millis(100)).await;
}
anyhow::Ok(())
})
.detach_and_log_err(cx);
anyhow::Ok(())
})
.detach_and_log_err(cx);
}
this
})?;

1
crates/context_servers/Cargo.toml
View file

@ -21,6 +21,7 @@ gpui.workspace = true
log.workspace = true
parking_lot.workspace = true
postage.workspace = true
project.workspace = true
schemars.workspace = true
serde.workspace = true
serde_json.workspace = true

2
crates/context_servers/src/client.rs
View file

@ -53,7 +53,7 @@ pub struct Client {
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ContextServerId(pub String);
pub struct ContextServerId(pub Arc<str>);
fn is_null_value<T: Serialize>(value: &T) -> bool {
if let Ok(Value::Null) = serde_json::to_value(value) {

57
crates/context_servers/src/manager.rs
View file

@ -18,7 +18,7 @@ use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use anyhow::Result;
use anyhow::{bail, Result};
use async_trait::async_trait;
use collections::{HashMap, HashSet};
use futures::{Future, FutureExt};
@ -36,19 +36,25 @@ use crate::{
#[derive(Deserialize, Serialize, Default, Clone, PartialEq, Eq, JsonSchema, Debug)]
pub struct ContextServerSettings {
pub servers: Vec<ServerConfig>,
#[serde(default)]
pub context_servers: HashMap<Arc<str>, ServerConfig>,
}
#[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema, Debug, Default)]
pub struct ServerConfig {
pub command: Option<ServerCommand>,
pub settings: Option<serde_json::Value>,
}
#[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema, Debug)]
pub struct ServerConfig {
pub id: String,
pub executable: String,
pub struct ServerCommand {
pub path: String,
pub args: Vec<String>,
pub env: Option<HashMap<String, String>>,
}
impl Settings for ContextServerSettings {
const KEY: Option<&'static str> = Some("experimental.context_servers");
const KEY: Option<&'static str> = None;
type FileContent = Self;
@ -79,9 +85,9 @@ pub struct NativeContextServer {
}
impl NativeContextServer {
pub fn new(config: Arc<ServerConfig>) -> Self {
pub fn new(id: Arc<str>, config: Arc<ServerConfig>) -> Self {
Self {
id: config.id.clone().into(),
id,
config,
client: RwLock::new(None),
}
@ -107,13 +113,16 @@ impl ContextServer for NativeContextServer {
cx: &'a AsyncAppContext,
) -> Pin<Box<dyn 'a + Future<Output = Result<()>>>> {
async move {
log::info!("starting context server {}", self.config.id,);
log::info!("starting context server {}", self.id);
let Some(command) = &self.config.command else {
bail!("no command specified for server {}", self.id);
};
let client = Client::new(
client::ContextServerId(self.config.id.clone()),
client::ContextServerId(self.id.clone()),
client::ModelContextServerBinary {
executable: Path::new(&self.config.executable).to_path_buf(),
args: self.config.args.clone(),
env: self.config.env.clone(),
executable: Path::new(&command.path).to_path_buf(),
args: command.args.clone(),
env: command.env.clone(),
},
cx.clone(),
)?;
@ -127,7 +136,7 @@ impl ContextServer for NativeContextServer {
log::debug!(
"context server {} initialized: {:?}",
self.config.id,
self.id,
initialized_protocol.initialize,
);
@ -242,7 +251,7 @@ impl ContextServerManager {
if let Some(server) = this.update(&mut cx, |this, _cx| this.servers.remove(&id))? {
server.stop()?;
let config = server.config();
let new_server = Arc::new(NativeContextServer::new(config));
let new_server = Arc::new(NativeContextServer::new(id.clone(), config));
new_server.clone().start(&cx).await?;
this.update(&mut cx, |this, cx| {
this.servers.insert(id.clone(), new_server);
@ -270,15 +279,15 @@ impl ContextServerManager {
.collect::<HashMap<_, _>>();
let new_servers = settings
.servers
.context_servers
.iter()
.map(|config| (config.id.clone(), config.clone()))
.map(|(id, config)| (id.clone(), config.clone()))
.collect::<HashMap<_, _>>();
let servers_to_add = new_servers
.values()
.filter(|config| !current_servers.contains_key(config.id.as_str()))
.cloned()
.iter()
.filter(|(id, _)| !current_servers.contains_key(id.as_ref()))
.map(|(id, config)| (id.clone(), config.clone()))
.collect::<Vec<_>>();
let servers_to_remove = current_servers
@ -288,9 +297,11 @@ impl ContextServerManager {
.collect::<Vec<_>>();
log::trace!("servers_to_add={:?}", servers_to_add);
for config in servers_to_add {
let server = Arc::new(NativeContextServer::new(Arc::new(config)));
self.add_server(server, cx).detach_and_log_err(cx);
for (id, config) in servers_to_add {
if config.command.is_some() {
let server = Arc::new(NativeContextServer::new(id, Arc::new(config)));
self.add_server(server, cx).detach_and_log_err(cx);
}
}
for id in servers_to_remove {

12
crates/context_servers/src/registry.rs
View file

@ -2,14 +2,18 @@ use std::sync::Arc;
use anyhow::Result;
use collections::HashMap;
use gpui::{AppContext, AsyncAppContext, ReadGlobal};
use gpui::{Global, Task};
use gpui::{AppContext, AsyncAppContext, Global, Model, ReadGlobal, Task};
use parking_lot::RwLock;
use project::Project;
use crate::ContextServer;
pub type ContextServerFactory =
Arc<dyn Fn(&AsyncAppContext) -> Task<Result<Arc<dyn ContextServer>>> + Send + Sync + 'static>;
pub type ContextServerFactory = Arc<
dyn Fn(Model<Project>, &AsyncAppContext) -> Task<Result<Arc<dyn ContextServer>>>
+ Send
+ Sync
+ 'static,
>;
#[derive(Default)]
struct GlobalContextServerFactoryRegistry(Arc<ContextServerFactoryRegistry>);

15
crates/extension_api/src/extension_api.rs
View file

@ -28,7 +28,7 @@ pub use wit::{
SlashCommand, SlashCommandArgumentCompletion, SlashCommandOutput, SlashCommandOutputSection,
},
CodeLabel, CodeLabelSpan, CodeLabelSpanLiteral, Command, DownloadedFileType, EnvVars,
KeyValueStore, LanguageServerInstallationStatus, Range, Worktree,
KeyValueStore, LanguageServerInstallationStatus, Project, Range, Worktree,
};
// Undocumented WIT re-exports.
@ -130,7 +130,11 @@ pub trait Extension: Send + Sync {
}
/// Returns the command used to start a context server.
fn context_server_command(&mut self, _context_server_id: &ContextServerId) -> Result<Command> {
fn context_server_command(
&mut self,
_context_server_id: &ContextServerId,
_project: &Project,
) -> Result<Command> {
Err("`context_server_command` not implemented".to_string())
}
@ -275,9 +279,12 @@ impl wit::Guest for Component {
extension().run_slash_command(command, args, worktree)
}
fn context_server_command(context_server_id: String) -> Result<wit::Command> {
fn context_server_command(
context_server_id: String,
project: &Project,
) -> Result<wit::Command> {
let context_server_id = ContextServerId(context_server_id);
extension().context_server_command(&context_server_id)
extension().context_server_command(&context_server_id, project)
}
fn suggest_docs_packages(provider: String) -> Result<Vec<String>, String> {

54
crates/extension_api/src/settings.rs
View file

@ -1,34 +1,56 @@
//! Provides access to Zed settings.
#[path = "../wit/since_v0.1.0/settings.rs"]
#[path = "../wit/since_v0.2.0/settings.rs"]
mod types;
use crate::{wit, Result, SettingsLocation, Worktree};
use crate::{wit, Project, Result, SettingsLocation, Worktree};
use serde_json;
pub use types::*;
impl LanguageSettings {
/// Returns the [`LanguageSettings`] for the given language.
pub fn for_worktree(language: Option<&str>, worktree: &Worktree) -> Result<Self> {
let location = SettingsLocation {
worktree_id: worktree.id(),
path: worktree.root_path(),
};
let settings_json = wit::get_settings(Some(&location), "language", language)?;
let settings: Self = serde_json::from_str(&settings_json).map_err(|err| err.to_string())?;
Ok(settings)
get_settings("language", language, Some(worktree.id()))
}
}
impl LspSettings {
/// Returns the [`LspSettings`] for the given language server.
pub fn for_worktree(language_server_name: &str, worktree: &Worktree) -> Result<Self> {
let location = SettingsLocation {
worktree_id: worktree.id(),
path: worktree.root_path(),
};
let settings_json = wit::get_settings(Some(&location), "lsp", Some(language_server_name))?;
let settings: Self = serde_json::from_str(&settings_json).map_err(|err| err.to_string())?;
Ok(settings)
get_settings("lsp", Some(language_server_name), Some(worktree.id()))
}
}
impl ContextServerSettings {
/// Returns the [`ContextServerSettings`] for the given context server.
pub fn for_project(context_server_id: &str, project: &Project) -> Result<Self> {
let global_setting: Self = get_settings("context_servers", Some(context_server_id), None)?;
for worktree_id in project.worktree_ids() {
let settings = get_settings(
"context_servers",
Some(context_server_id),
Some(worktree_id),
)?;
if settings != global_setting {
return Ok(settings);
}
}
Ok(global_setting)
}
}
fn get_settings<T: serde::de::DeserializeOwned>(
settings_type: &str,
settings_name: Option<&str>,
worktree_id: Option<u64>,
) -> Result<T> {
let location = worktree_id.map(|worktree_id| SettingsLocation {
worktree_id,
path: String::new(),
});
let settings_json = wit::get_settings(location.as_ref(), settings_type, settings_name)?;
let settings: T = serde_json::from_str(&settings_json).map_err(|err| err.to_string())?;
Ok(settings)
}

8
crates/extension_api/wit/since_v0.2.0/extension.wit
View file

@ -83,6 +83,12 @@ world extension {
shell-env: func() -> env-vars;
}
/// A Zed project.
resource project {
/// Returns the IDs of all of the worktrees in this project.
worktree-ids: func() -> list<u64>;
}
/// A key-value store.
resource key-value-store {
/// Inserts an entry under the specified key.
@ -136,7 +142,7 @@ world extension {
export run-slash-command: func(command: slash-command, args: list<string>, worktree: option<borrow<worktree>>) -> result<slash-command-output, string>;
/// Returns the command used to start up a context server.
export context-server-command: func(context-server-id: string) -> result<command, string>;
export context-server-command: func(context-server-id: string, project: borrow<project>) -> result<command, string>;
/// Returns a list of packages as suggestions to be included in the `/docs`
/// search results.

29
crates/extension_api/wit/since_v0.2.0/settings.rs
View file

@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use std::num::NonZeroU32;
use std::{collections::HashMap, num::NonZeroU32};
/// The settings for a particular language.
#[derive(Debug, Serialize, Deserialize)]
@ -12,18 +12,29 @@ pub struct LanguageSettings {
#[derive(Default, Debug, Serialize, Deserialize)]
pub struct LspSettings {
/// The settings for the language server binary.
pub binary: Option<BinarySettings>,
pub binary: Option<CommandSettings>,
/// The initialization options to pass to the language server.
pub initialization_options: Option<serde_json::Value>,
/// The settings to pass to language server.
pub settings: Option<serde_json::Value>,
}
/// The settings for a language server binary.
#[derive(Debug, Serialize, Deserialize)]
pub struct BinarySettings {
/// The path to the binary.
pub path: Option<String>,
/// The arguments to pass to the binary.
pub arguments: Option<Vec<String>>,
/// The settings for a particular context server.
#[derive(Default, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ContextServerSettings {
/// The settings for the context server binary.
pub command: Option<CommandSettings>,
/// The settings to pass to the context server.
pub settings: Option<serde_json::Value>,
}
/// The settings for a command.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct CommandSettings {
/// The path to the command.
pub path: Option<String>,
/// The arguments to pass to the command.
pub arguments: Option<Vec<String>>,
/// The environment variables.
pub env: Option<HashMap<String, String>>,
}

1
crates/extension_host/Cargo.toml
View file

@ -22,6 +22,7 @@ async-tar.workspace = true
async-trait.workspace = true
client.workspace = true
collections.workspace = true
context_servers.workspace = true
extension.workspace = true
fs.workspace = true
futures.workspace = true

2
crates/extension_host/src/wasm_host.rs
View file

@ -27,7 +27,7 @@ use wasmtime::{
};
use wasmtime_wasi as wasi;
use wit::Extension;
pub use wit::SlashCommand;
pub use wit::{ExtensionProject, SlashCommand};
pub struct WasmHost {
engine: Engine,

6
crates/extension_host/src/wasm_host/wit.rs
View file

@ -4,7 +4,6 @@ mod since_v0_0_6;
mod since_v0_1_0;
mod since_v0_2_0;
use lsp::LanguageServerName;
// use indexed_docs::IndexedDocsDatabase;
use release_channel::ReleaseChannel;
use since_v0_2_0 as latest;
@ -27,7 +26,7 @@ pub use latest::{
Completion, CompletionKind, CompletionLabelDetails, InsertTextFormat, Symbol, SymbolKind,
},
zed::extension::slash_command::{SlashCommandArgumentCompletion, SlashCommandOutput},
CodeLabel, CodeLabelSpan, Command, Range, SlashCommand,
CodeLabel, CodeLabelSpan, Command, ExtensionProject, Range, SlashCommand,
};
pub use since_v0_0_4::LanguageServerConfig;
@ -389,10 +388,11 @@ impl Extension {
&self,
store: &mut Store<WasmState>,
context_server_id: Arc<str>,
project: Resource<ExtensionProject>,
) -> Result<Result<Command, String>> {
match self {
Extension::V020(ext) => {
ext.call_context_server_command(store, &context_server_id)
ext.call_context_server_command(store, &context_server_id, project)
.await
}
Extension::V001(_) | Extension::V004(_) | Extension::V006(_) | Extension::V010(_) => {

43
crates/extension_host/src/wasm_host/wit/since_v0_2_0.rs
View file

@ -6,6 +6,7 @@ use anyhow::{anyhow, bail, Context, Result};
use async_compression::futures::bufread::GzipDecoder;
use async_tar::Archive;
use async_trait::async_trait;
use context_servers::manager::ContextServerSettings;
use futures::{io::BufReader, FutureExt as _};
use futures::{lock::Mutex, AsyncReadExt};
use language::{
@ -31,6 +32,7 @@ wasmtime::component::bindgen!({
path: "../extension_api/wit/since_v0.2.0",
with: {
"worktree": ExtensionWorktree,
"project": ExtensionProject,
"key-value-store": ExtensionKeyValueStore,
"zed:extension/http-client/http-response-stream": ExtensionHttpResponseStream
},
@ -46,6 +48,10 @@ pub type ExtensionWorktree = Arc<dyn LspAdapterDelegate>;
pub type ExtensionKeyValueStore = Arc<dyn DocsDatabase>;
pub type ExtensionHttpResponseStream = Arc<Mutex<::http_client::Response<AsyncBody>>>;
pub struct ExtensionProject {
pub worktree_ids: Vec<u64>,
}
pub fn linker() -> &'static Linker<WasmState> {
static LINKER: OnceLock<Linker<WasmState>> = OnceLock::new();
LINKER.get_or_init(|| super::new_linker(Extension::add_to_linker))
@ -69,6 +75,22 @@ impl HostKeyValueStore for WasmState {
}
}
#[async_trait]
impl HostProject for WasmState {
async fn worktree_ids(
&mut self,
project: Resource<ExtensionProject>,
) -> wasmtime::Result<Vec<u64>> {
let project = self.table.get(&project)?;
Ok(project.worktree_ids.clone())
}
fn drop(&mut self, _project: Resource<Project>) -> Result<()> {
// We only ever hand out borrows of projects.
Ok(())
}
}
#[async_trait]
impl HostWorktree for WasmState {
async fn id(
@ -421,14 +443,33 @@ impl ExtensionImports for WasmState {
.cloned()
.unwrap_or_default();
Ok(serde_json::to_string(&settings::LspSettings {
binary: settings.binary.map(|binary| settings::BinarySettings {
binary: settings.binary.map(|binary| settings::CommandSettings {
path: binary.path,
arguments: binary.arguments,
env: None,
}),
settings: settings.settings,
initialization_options: settings.initialization_options,
})?)
}
"context_servers" => {
let settings = key
.and_then(|key| {
ContextServerSettings::get(location, cx)
.context_servers
.get(key.as_str())
})
.cloned()
.unwrap_or_default();
Ok(serde_json::to_string(&settings::ContextServerSettings {
command: settings.command.map(|command| settings::CommandSettings {
path: Some(command.path),
arguments: Some(command.args),
env: command.env.map(|env| env.into_iter().collect()),
}),
settings: settings.settings,
})?)
}
_ => {
bail!("Unknown settings category: {}", category);
}

39
crates/extensions_ui/src/extension_context_server.rs
View file

@ -3,12 +3,14 @@ use std::sync::Arc;
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use context_servers::manager::{NativeContextServer, ServerConfig};
use context_servers::manager::{NativeContextServer, ServerCommand, ServerConfig};
use context_servers::protocol::InitializedContextServerProtocol;
use context_servers::ContextServer;
use extension_host::wasm_host::{WasmExtension, WasmHost};
use extension_host::wasm_host::{ExtensionProject, WasmExtension, WasmHost};
use futures::{Future, FutureExt};
use gpui::AsyncAppContext;
use gpui::{AsyncAppContext, Model};
use project::Project;
use wasmtime_wasi::WasiView as _;
pub struct ExtensionContextServer {
#[allow(unused)]
@ -20,14 +22,27 @@ pub struct ExtensionContextServer {
}
impl ExtensionContextServer {
pub async fn new(extension: WasmExtension, host: Arc<WasmHost>, id: Arc<str>) -> Result<Self> {
pub async fn new(
extension: WasmExtension,
host: Arc<WasmHost>,
id: Arc<str>,
project: Model<Project>,
mut cx: AsyncAppContext,
) -> Result<Self> {
let extension_project = project.update(&mut cx, |project, cx| ExtensionProject {
worktree_ids: project
.visible_worktrees(cx)
.map(|worktree| worktree.read(cx).id().to_proto())
.collect(),
})?;
let command = extension
.call({
let id = id.clone();
|extension, store| {
async move {
let project = store.data_mut().table().push(extension_project)?;
let command = extension
.call_context_server_command(store, id.clone())
.call_context_server_command(store, id.clone(), project)
.await?
.map_err(|e| anyhow!("{}", e))?;
anyhow::Ok(command)
@ -38,17 +53,19 @@ impl ExtensionContextServer {
.await?;
let config = Arc::new(ServerConfig {
id: id.to_string(),
executable: command.command,
args: command.args,
env: Some(command.env.into_iter().collect()),
settings: None,
command: Some(ServerCommand {
path: command.command,
args: command.args,
env: Some(command.env.into_iter().collect()),
}),
});
anyhow::Ok(Self {
extension,
host,
id,
context_server: Arc::new(NativeContextServer::new(config)),
id: id.clone(),
context_server: Arc::new(NativeContextServer::new(id, config)),
})
}
}

8
crates/extensions_ui/src/extension_registration_hooks.rs
View file

@ -84,14 +84,14 @@ impl extension_host::ExtensionRegistrationHooks for ConcreteExtensionRegistratio
.register_server_factory(
id.clone(),
Arc::new({
move |cx| {
move |project, cx| {
let id = id.clone();
let extension = extension.clone();
let host = host.clone();
cx.spawn(|_cx| async move {
cx.spawn(|cx| async move {
let context_server =
ExtensionContextServer::new(extension, host, id).await?;
ExtensionContextServer::new(extension, host, id, project, cx)
.await?;
anyhow::Ok(Arc::new(context_server) as _)
})
}