oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Put context servers behind a trait (#20432)

Browse Source 
This PR puts context servers behind the `ContextServer` trait to allow
us to provide context servers from an extension.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers
2024-11-08 13:36:41 -05:00
committed by GitHub
parent 01503511ad
commit 09c599385a
6 changed files with 102 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/context_servers/src/manager.rs
+87 -50
View File
@@ -15,9 +15,13 @@
//! and react to changes in settings.
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use collections::{HashMap, HashSet};
use futures::{Future, FutureExt};
use gpui::{AsyncAppContext, EventEmitter, ModelContext, Task};
use log;
use parking_lot::RwLock;
@@ -56,51 +60,84 @@ impl Settings for ContextServerSettings {
}
}
pub struct ContextServer {
pub id: String,
pub config: ServerConfig,
#[async_trait(?Send)]
pub trait ContextServer: Send + Sync + 'static {
fn id(&self) -> Arc<str>;
fn config(&self) -> Arc<ServerConfig>;
fn client(&self) -> Option<Arc<crate::protocol::InitializedContextServerProtocol>>;
fn start<'a>(
self: Arc<Self>,
cx: &'a AsyncAppContext,
) -> Pin<Box<dyn 'a + Future<Output = Result<()>>>>;
fn stop(&self) -> Result<()>;
}
pub struct NativeContextServer {
pub id: Arc<str>,
pub config: Arc<ServerConfig>,
pub client: RwLock<Option<Arc<crate::protocol::InitializedContextServerProtocol>>>,
}
impl ContextServer {
fn new(config: ServerConfig) -> Self {
impl NativeContextServer {
fn new(config: Arc<ServerConfig>) -> Self {
Self {
id: config.id.clone(),
id: config.id.clone().into(),
config,
client: RwLock::new(None),
}
}
}
async fn start(&self, cx: &AsyncAppContext) -> anyhow::Result<()> {
log::info!("starting context server {}", self.config.id,);
let client = Client::new(
client::ContextServerId(self.config.id.clone()),
client::ModelContextServerBinary {
executable: Path::new(&self.config.executable).to_path_buf(),
args: self.config.args.clone(),
env: self.config.env.clone(),
},
cx.clone(),
)?;
let protocol = crate::protocol::ModelContextProtocol::new(client);
let client_info = types::Implementation {
name: "Zed".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
};
let initialized_protocol = protocol.initialize(client_info).await?;
log::debug!(
"context server {} initialized: {:?}",
self.config.id,
initialized_protocol.initialize,
);
*self.client.write() = Some(Arc::new(initialized_protocol));
Ok(())
#[async_trait(?Send)]
impl ContextServer for NativeContextServer {
fn id(&self) -> Arc<str> {
self.id.clone()
}
async fn stop(&self) -> anyhow::Result<()> {
fn config(&self) -> Arc<ServerConfig> {
self.config.clone()
}
fn client(&self) -> Option<Arc<crate::protocol::InitializedContextServerProtocol>> {
self.client.read().clone()
}
fn start<'a>(
self: Arc<Self>,
cx: &'a AsyncAppContext,
) -> Pin<Box<dyn 'a + Future<Output = Result<()>>>> {
async move {
log::info!("starting context server {}", self.config.id,);
let client = Client::new(
client::ContextServerId(self.config.id.clone()),
client::ModelContextServerBinary {
executable: Path::new(&self.config.executable).to_path_buf(),
args: self.config.args.clone(),
env: self.config.env.clone(),
},
cx.clone(),
)?;
let protocol = crate::protocol::ModelContextProtocol::new(client);
let client_info = types::Implementation {
name: "Zed".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
};
let initialized_protocol = protocol.initialize(client_info).await?;
log::debug!(
"context server {} initialized: {:?}",
self.config.id,
initialized_protocol.initialize,
);
*self.client.write() = Some(Arc::new(initialized_protocol));
Ok(())
}
.boxed_local()
}
fn stop(&self) -> Result<()> {
let mut client = self.client.write();
if let Some(protocol) = client.take() {
drop(protocol);
@@ -114,7 +151,7 @@ impl ContextServer {
/// must go through the `GlobalContextServerManager` which holds
/// a model to the ContextServerManager.
pub struct ContextServerManager {
servers: HashMap<String, Arc<ContextServer>>,
servers: HashMap<String, Arc<dyn ContextServer>>,
pending_servers: HashSet<String>,
}
@@ -141,7 +178,7 @@ impl ContextServerManager {
pub fn add_server(
&mut self,
config: ServerConfig,
config: Arc<ServerConfig>,
cx: &ModelContext<Self>,
) -> Task<anyhow::Result<()>> {
let server_id = config.id.clone();
@@ -153,8 +190,8 @@ impl ContextServerManager {
let task = {
let server_id = server_id.clone();
cx.spawn(|this, mut cx| async move {
let server = Arc::new(ContextServer::new(config));
server.start(&cx).await?;
let server = Arc::new(NativeContextServer::new(config));
server.clone().start(&cx).await?;
this.update(&mut cx, |this, cx| {
this.servers.insert(server_id.clone(), server);
this.pending_servers.remove(&server_id);
@@ -170,7 +207,7 @@ impl ContextServerManager {
task
}
pub fn get_server(&self, id: &str) -> Option<Arc<ContextServer>> {
pub fn get_server(&self, id: &str) -> Option<Arc<dyn ContextServer>> {
self.servers.get(id).cloned()
}
@@ -178,7 +215,7 @@ impl ContextServerManager {
let id = id.to_string();
cx.spawn(|this, mut cx| async move {
if let Some(server) = this.update(&mut cx, |this, _cx| this.servers.remove(&id))? {
server.stop().await?;
server.stop()?;
}
this.update(&mut cx, |this, cx| {
this.pending_servers.remove(&id);
@@ -192,16 +229,16 @@ impl ContextServerManager {
pub fn restart_server(
&mut self,
id: &str,
id: &Arc<str>,
cx: &mut ModelContext<Self>,
) -> Task<anyhow::Result<()>> {
let id = id.to_string();
cx.spawn(|this, mut cx| async move {
if let Some(server) = this.update(&mut cx, |this, _cx| this.servers.remove(&id))? {
server.stop().await?;
let config = server.config.clone();
let new_server = Arc::new(ContextServer::new(config));
new_server.start(&cx).await?;
server.stop()?;
let config = server.config();
let new_server = Arc::new(NativeContextServer::new(config));
new_server.clone().start(&cx).await?;
this.update(&mut cx, |this, cx| {
this.servers.insert(id.clone(), new_server);
cx.emit(Event::ServerStopped {
@@ -216,7 +253,7 @@ impl ContextServerManager {
})
}
pub fn servers(&self) -> Vec<Arc<ContextServer>> {
pub fn servers(&self) -> Vec<Arc<dyn ContextServer>> {
self.servers.values().cloned().collect()
}
@@ -224,7 +261,7 @@ impl ContextServerManager {
let current_servers = self
.servers()
.into_iter()
.map(|server| (server.id.clone(), server.config.clone()))
.map(|server| (server.id(), server.config()))
.collect::<HashMap<_, _>>();
let new_servers = settings
@@ -235,19 +272,19 @@ impl ContextServerManager {
let servers_to_add = new_servers
.values()
.filter(|config| !current_servers.contains_key(&config.id))
.filter(|config| !current_servers.contains_key(config.id.as_str()))
.cloned()
.collect::<Vec<_>>();
let servers_to_remove = current_servers
.keys()
.filter(|id| !new_servers.contains_key(*id))
.filter(|id| !new_servers.contains_key(id.as_ref()))
.cloned()
.collect::<Vec<_>>();
log::trace!("servers_to_add={:?}", servers_to_add);
for config in servers_to_add {
self.add_server(config, cx).detach_and_log_err(cx);
self.add_server(Arc::new(config), cx).detach_and_log_err(cx);
}
for id in servers_to_remove {