use std::rc::Rc; use acp_thread::{AgentConnection, LoadError};use agent_servers::AcpConnection;use agent_servers::{AgentServer, AgentServerDelegate};use anyhow::Result;use collections::HashMap;use futures::{FutureExt, future::Shared};use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Subscription, Task}; use project::{AgentServerStore, AgentServersUpdated, Project};use watch::Receiver; use crate::Agent; pub enum AgentConnectionEntry { Connecting { connect_task: Shared<Task<Result<AgentConnectedState, LoadError>>>, }, Connected(AgentConnectedState), Error { error: LoadError, },} #[derive(Clone)]pub struct AgentConnectedState { pub connection: Rc<dyn AgentConnection>,} #[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum AgentConnectionStatus { Disconnected, Connecting, Connected,} impl AgentConnectionEntry { pub fn wait_for_connection(&self) -> Shared<Task<Result<AgentConnectedState, LoadError>>> { match self { AgentConnectionEntry::Connecting { connect_task } => connect_task.clone(), AgentConnectionEntry::Connected(state) => Task::ready(Ok(state.clone())).shared(), AgentConnectionEntry::Error { error } => Task::ready(Err(error.clone())).shared(), } } pub fn status(&self) -> AgentConnectionStatus { match self { AgentConnectionEntry::Connecting { .. } => AgentConnectionStatus::Connecting, AgentConnectionEntry::Connected(_) => AgentConnectionStatus::Connected, AgentConnectionEntry::Error { .. } => AgentConnectionStatus::Disconnected, } }} pub enum AgentConnectionEntryEvent { NewVersionAvailable(SharedString), LoadingStatusChanged(Option<SharedString>),} impl EventEmitter<AgentConnectionEntryEvent> for AgentConnectionEntry {} #[derive(Clone)]pub struct ActiveAcpConnection { pub agent_id: project::AgentId, pub connection: Rc<AcpConnection>,} pub struct AgentConnectionStore { project: Entity<Project>, entries: HashMap<Agent, Entity<AgentConnectionEntry>>, _subscriptions: Vec<Subscription>,} impl AgentConnectionStore { pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self { let agent_server_store = project.read(cx).agent_server_store().clone(); let subscription = cx.subscribe(&agent_server_store, Self::handle_agent_servers_updated); Self { project, entries: HashMap::default(), _subscriptions: vec![subscription], } } pub fn project(&self) -> &Entity<Project> { &self.project } pub fn entry(&self, key: &Agent) -> Option<&Entity<AgentConnectionEntry>> { self.entries.get(key) } pub fn connection_status(&self, key: &Agent, cx: &App) -> AgentConnectionStatus { self.entries .get(key) .map(|entry| entry.read(cx).status()) .unwrap_or(AgentConnectionStatus::Disconnected) } pub fn agent_version(&self, key: &Agent, cx: &App) -> Option<SharedString> { match self.entries.get(key)?.read(cx) { AgentConnectionEntry::Connected(state) => state.connection.agent_version(), AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => None, } } pub fn active_acp_connections(&self, cx: &App) -> Vec<ActiveAcpConnection> { self.entries .values() .filter_map(|entry| match entry.read(cx) { AgentConnectionEntry::Connected(state) => state .connection .clone() .downcast::<AcpConnection>() .map(|connection| ActiveAcpConnection { agent_id: state.connection.agent_id(), connection, }), AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => { None } }) .collect() } pub fn restart_connection( &mut self, key: Agent, server: Rc<dyn AgentServer>, cx: &mut Context<Self>, ) -> Entity<AgentConnectionEntry> { if let Some(entry) = self.entries.get(&key) { if matches!(entry.read(cx), AgentConnectionEntry::Connecting { .. }) { return entry.clone(); } } self.entries.remove(&key); self.request_connection(key, server, cx) } pub fn request_connection( &mut self, key: Agent, server: Rc<dyn AgentServer>, cx: &mut Context<Self>, ) -> Entity<AgentConnectionEntry> { if let Some(entry) = self.entries.get(&key) { return entry.clone(); } let (mut new_version_rx, mut loading_status_rx, connect_task) = self.start_connection(server, cx); let connect_task = connect_task.shared(); let entry = cx.new(|_cx| AgentConnectionEntry::Connecting { connect_task: connect_task.clone(), }); self.entries.insert(key.clone(), entry.clone()); cx.notify(); cx.spawn({ let key = key.clone(); let entry = entry.downgrade(); async move |this, cx| match connect_task.await { Ok(connected_state) => { this.update(cx, move |this, cx| { if this.entries.get(&key) != entry.upgrade().as_ref() { return; } entry .update(cx, move |entry, cx| { if let AgentConnectionEntry::Connecting { .. } = entry { *entry = AgentConnectionEntry::Connected(connected_state); cx.notify(); } }) .ok(); cx.notify(); }) .ok(); } Err(error) => { this.update(cx, move |this, cx| { if this.entries.get(&key) != entry.upgrade().as_ref() { return; } entry .update(cx, move |entry, cx| { if let AgentConnectionEntry::Connecting { .. } = entry { *entry = AgentConnectionEntry::Error { error }; cx.notify(); } }) .ok(); this.entries.remove(&key); cx.notify(); }) .ok(); } } }) .detach(); cx.spawn({ let key = key.clone(); let entry = entry.downgrade(); async move |this, cx| { while let Ok(version) = new_version_rx.recv().await { let Some(version) = version else { continue; }; this.update(cx, move |this, cx| { if this.entries.get(&key) != entry.upgrade().as_ref() { return; } entry .update(cx, move |_entry, cx| { cx.emit(AgentConnectionEntryEvent::NewVersionAvailable( version.into(), )); }) .ok(); this.entries.remove(&key); cx.notify(); }) .ok(); break; } } }) .detach(); cx.spawn({ let entry = entry.downgrade(); async move |this, cx| { while let Ok(status) = loading_status_rx.recv().await { let status = status.map(SharedString::from); let key = key.clone(); let entry = entry.clone(); this.update(cx, move |this, cx| { if this.entries.get(&key) != entry.upgrade().as_ref() { return; } entry .update(cx, move |_entry, cx| { cx.emit(AgentConnectionEntryEvent::LoadingStatusChanged(status)); }) .ok(); cx.notify(); }) .ok(); } } }) .detach(); entry } fn handle_agent_servers_updated( &mut self, store: Entity<AgentServerStore>, _: &AgentServersUpdated, cx: &mut Context<Self>, ) { let store = store.read(cx); self.entries.retain(|key, _| match key { Agent::NativeAgent => true, Agent::Custom { id } => store.external_agents.contains_key(id), #[cfg(any(test, feature = "test-support"))] Agent::Stub => true, }); cx.notify(); } fn start_connection( &self, server: Rc<dyn AgentServer>, cx: &mut Context<Self>, ) -> ( Receiver<Option<String>>, Receiver<Option<String>>, Task<Result<AgentConnectedState, LoadError>>, ) { let (new_version_tx, new_version_rx) = watch::channel::<Option<String>>(None); let (loading_status_tx, loading_status_rx) = watch::channel::<Option<String>>(None); let agent_server_store = self.project.read(cx).agent_server_store().clone(); let delegate = AgentServerDelegate::new( agent_server_store, Some(new_version_tx), Some(loading_status_tx), ); let connect_task = server.connect(delegate, self.project.clone(), cx); let connect_task = cx.spawn(async move |_this, _cx| match connect_task.await { Ok(connection) => Ok(AgentConnectedState { connection }), Err(err) => match err.downcast::<LoadError>() { Ok(load_error) => Err(load_error), Err(err) => Err(LoadError::Other(SharedString::from(err.to_string()))), }, }); (new_version_rx, loading_status_rx, connect_task) }}
use std::rc::Rc;
use acp_thread::{AgentConnection, LoadError};
use agent_servers::AcpConnection;
use agent_servers::{AgentServer, AgentServerDelegate};
use anyhow::Result;
use collections::HashMap;
use futures::{FutureExt, future::Shared};
use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Subscription, Task};
use project::{AgentServerStore, AgentServersUpdated, Project};
use watch::Receiver;
use crate::Agent;
pub enum AgentConnectionEntry {
Connecting {
connect_task: Shared<Task<Result<AgentConnectedState, LoadError>>>,
},
Connected(AgentConnectedState),
Error {
error: LoadError,
},
}
#[derive(Clone)]
pub struct AgentConnectedState {
pub connection: Rc<dyn AgentConnection>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AgentConnectionStatus {
Disconnected,
Connecting,
Connected,
}
impl AgentConnectionEntry {
pub fn wait_for_connection(&self) -> Shared<Task<Result<AgentConnectedState, LoadError>>> {
match self {
AgentConnectionEntry::Connecting { connect_task } => connect_task.clone(),
AgentConnectionEntry::Connected(state) => Task::ready(Ok(state.clone())).shared(),
AgentConnectionEntry::Error { error } => Task::ready(Err(error.clone())).shared(),
}
}
pub fn status(&self) -> AgentConnectionStatus {
match self {
AgentConnectionEntry::Connecting { .. } => AgentConnectionStatus::Connecting,
AgentConnectionEntry::Connected(_) => AgentConnectionStatus::Connected,
AgentConnectionEntry::Error { .. } => AgentConnectionStatus::Disconnected,
}
}
}
pub enum AgentConnectionEntryEvent {
NewVersionAvailable(SharedString),
LoadingStatusChanged(Option<SharedString>),
}
impl EventEmitter<AgentConnectionEntryEvent> for AgentConnectionEntry {}
#[derive(Clone)]
pub struct ActiveAcpConnection {
pub agent_id: project::AgentId,
pub connection: Rc<AcpConnection>,
}
pub struct AgentConnectionStore {
project: Entity<Project>,
entries: HashMap<Agent, Entity<AgentConnectionEntry>>,
_subscriptions: Vec<Subscription>,
}
impl AgentConnectionStore {
pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self {
let agent_server_store = project.read(cx).agent_server_store().clone();
let subscription = cx.subscribe(&agent_server_store, Self::handle_agent_servers_updated);
Self {
project,
entries: HashMap::default(),
_subscriptions: vec![subscription],
}
}
pub fn project(&self) -> &Entity<Project> {
&self.project
}
pub fn entry(&self, key: &Agent) -> Option<&Entity<AgentConnectionEntry>> {
self.entries.get(key)
}
pub fn connection_status(&self, key: &Agent, cx: &App) -> AgentConnectionStatus {
self.entries
.get(key)
.map(|entry| entry.read(cx).status())
.unwrap_or(AgentConnectionStatus::Disconnected)
}
pub fn agent_version(&self, key: &Agent, cx: &App) -> Option<SharedString> {
match self.entries.get(key)?.read(cx) {
AgentConnectionEntry::Connected(state) => state.connection.agent_version(),
AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => None,
}
}
pub fn active_acp_connections(&self, cx: &App) -> Vec<ActiveAcpConnection> {
self.entries
.values()
.filter_map(|entry| match entry.read(cx) {
AgentConnectionEntry::Connected(state) => state
.connection
.clone()
.downcast::<AcpConnection>()
.map(|connection| ActiveAcpConnection {
agent_id: state.connection.agent_id(),
connection,
}),
AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => {
None
}
})
.collect()
}
pub fn restart_connection(
&mut self,
key: Agent,
server: Rc<dyn AgentServer>,
cx: &mut Context<Self>,
) -> Entity<AgentConnectionEntry> {
if let Some(entry) = self.entries.get(&key) {
if matches!(entry.read(cx), AgentConnectionEntry::Connecting { .. }) {
return entry.clone();
}
}
self.entries.remove(&key);
self.request_connection(key, server, cx)
}
pub fn request_connection(
&mut self,
key: Agent,
server: Rc<dyn AgentServer>,
cx: &mut Context<Self>,
) -> Entity<AgentConnectionEntry> {
if let Some(entry) = self.entries.get(&key) {
return entry.clone();
}
let (mut new_version_rx, mut loading_status_rx, connect_task) =
self.start_connection(server, cx);
let connect_task = connect_task.shared();
let entry = cx.new(|_cx| AgentConnectionEntry::Connecting {
connect_task: connect_task.clone(),
});
self.entries.insert(key.clone(), entry.clone());
cx.notify();
cx.spawn({
let key = key.clone();
let entry = entry.downgrade();
async move |this, cx| match connect_task.await {
Ok(connected_state) => {
this.update(cx, move |this, cx| {
if this.entries.get(&key) != entry.upgrade().as_ref() {
return;
}
entry
.update(cx, move |entry, cx| {
if let AgentConnectionEntry::Connecting { .. } = entry {
*entry = AgentConnectionEntry::Connected(connected_state);
cx.notify();
}
})
.ok();
cx.notify();
})
.ok();
}
Err(error) => {
this.update(cx, move |this, cx| {
if this.entries.get(&key) != entry.upgrade().as_ref() {
return;
}
entry
.update(cx, move |entry, cx| {
if let AgentConnectionEntry::Connecting { .. } = entry {
*entry = AgentConnectionEntry::Error { error };
cx.notify();
}
})
.ok();
this.entries.remove(&key);
cx.notify();
})
.ok();
}
}
})
.detach();
cx.spawn({
let key = key.clone();
let entry = entry.downgrade();
async move |this, cx| {
while let Ok(version) = new_version_rx.recv().await {
let Some(version) = version else {
continue;
};
this.update(cx, move |this, cx| {
if this.entries.get(&key) != entry.upgrade().as_ref() {
return;
}
entry
.update(cx, move |_entry, cx| {
cx.emit(AgentConnectionEntryEvent::NewVersionAvailable(
version.into(),
));
})
.ok();
this.entries.remove(&key);
cx.notify();
})
.ok();
break;
}
}
})
.detach();
cx.spawn({
let entry = entry.downgrade();
async move |this, cx| {
while let Ok(status) = loading_status_rx.recv().await {
let status = status.map(SharedString::from);
let key = key.clone();
let entry = entry.clone();
this.update(cx, move |this, cx| {
if this.entries.get(&key) != entry.upgrade().as_ref() {
return;
}
entry
.update(cx, move |_entry, cx| {
cx.emit(AgentConnectionEntryEvent::LoadingStatusChanged(status));
})
.ok();
cx.notify();
})
.ok();
}
}
})
.detach();
entry
}
fn handle_agent_servers_updated(
&mut self,
store: Entity<AgentServerStore>,
_: &AgentServersUpdated,
cx: &mut Context<Self>,
) {
let store = store.read(cx);
self.entries.retain(|key, _| match key {
Agent::NativeAgent => true,
Agent::Custom { id } => store.external_agents.contains_key(id),
#[cfg(any(test, feature = "test-support"))]
Agent::Stub => true,
});
cx.notify();
}
fn start_connection(
&self,
server: Rc<dyn AgentServer>,
cx: &mut Context<Self>,
) -> (
Receiver<Option<String>>,
Receiver<Option<String>>,
Task<Result<AgentConnectedState, LoadError>>,
) {
let (new_version_tx, new_version_rx) = watch::channel::<Option<String>>(None);
let (loading_status_tx, loading_status_rx) = watch::channel::<Option<String>>(None);
let agent_server_store = self.project.read(cx).agent_server_store().clone();
let delegate = AgentServerDelegate::new(
agent_server_store,
Some(new_version_tx),
Some(loading_status_tx),
);
let connect_task = server.connect(delegate, self.project.clone(), cx);
let connect_task = cx.spawn(async move |_this, _cx| match connect_task.await {
Ok(connection) => Ok(AgentConnectedState { connection }),
Err(err) => match err.downcast::<LoadError>() {
Ok(load_error) => Err(load_error),
Err(err) => Err(LoadError::Other(SharedString::from(err.to_string()))),
},
});
(new_version_rx, loading_status_rx, connect_task)
}
}