Consolidate logic for protobuf message handling between ssh and web socket clients (#17185)

This is a refactor to prepare for adding LSP support in SSH remote
projects.

Release Notes:

- N/A

---------

Co-authored-by: Mikayla <mikayla@zed.dev>
Co-authored-by: Conrad <conrad@zed.dev>

Cargo.lock

@@ -8512,6 +8512,8 @@ dependencies = [
  "anyhow",
  "collections",
  "futures 0.3.30",
+ "gpui",
+ "parking_lot",
  "prost",
  "prost-build",
  "serde",

crates/assistant/src/assistant.rs

@@ -209,7 +209,7 @@ pub fn init(
     })
     .detach();
 
-    context_store::init(&client);
+    context_store::init(&client.clone().into());
     prompt_library::init(cx);
     init_language_model_settings(cx);
     assistant_slash_command::init(cx);

crates/assistant/src/context_store.rs

@@ -2,6 +2,7 @@ use crate::{
     prompts::PromptBuilder, Context, ContextEvent, ContextId, ContextOperation, ContextVersion,
     SavedContext, SavedContextMetadata,
 };
+use ::proto::AnyProtoClient;
 use anyhow::{anyhow, Context as _, Result};
 use client::{proto, telemetry::Telemetry, Client, TypedEnvelope};
 use clock::ReplicaId;
@@ -25,7 +26,7 @@ use std::{
 };
 use util::{ResultExt, TryFutureExt};
 
-pub fn init(client: &Arc<Client>) {
+pub fn init(client: &AnyProtoClient) {
     client.add_model_message_handler(ContextStore::handle_advertise_contexts);
     client.add_model_request_handler(ContextStore::handle_open_context);
     client.add_model_request_handler(ContextStore::handle_create_context);

crates/channel/src/channel.rs

@@ -18,6 +18,6 @@ mod channel_store_tests;
 
 pub fn init(client: &Arc<Client>, user_store: Model<UserStore>, cx: &mut AppContext) {
     channel_store::init(client, user_store, cx);
-    channel_buffer::init(client);
+    channel_buffer::init(&client.clone().into());
-    channel_chat::init(client);
+    channel_chat::init(&client.clone().into());
 }

crates/channel/src/channel_buffer.rs

@@ -5,7 +5,7 @@ use collections::HashMap;
 use gpui::{AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, Task};
 use language::proto::serialize_version;
 use rpc::{
-    proto::{self, PeerId},
+    proto::{self, AnyProtoClient, PeerId},
     TypedEnvelope,
 };
 use std::{sync::Arc, time::Duration};
@@ -14,7 +14,7 @@ use util::ResultExt;
 
 pub const ACKNOWLEDGE_DEBOUNCE_INTERVAL: Duration = Duration::from_millis(250);
 
-pub(crate) fn init(client: &Arc<Client>) {
+pub(crate) fn init(client: &AnyProtoClient) {
     client.add_model_message_handler(ChannelBuffer::handle_update_channel_buffer);
     client.add_model_message_handler(ChannelBuffer::handle_update_channel_buffer_collaborators);
 }

crates/channel/src/channel_chat.rs

@@ -11,6 +11,7 @@ use gpui::{
     AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, Task, WeakModel,
 };
 use rand::prelude::*;
+use rpc::proto::AnyProtoClient;
 use std::{
     ops::{ControlFlow, Range},
     sync::Arc,
@@ -95,7 +96,7 @@ pub enum ChannelChatEvent {
 }
 
 impl EventEmitter<ChannelChatEvent> for ChannelChat {}
-pub fn init(client: &Arc<Client>) {
+pub fn init(client: &AnyProtoClient) {
     client.add_model_message_handler(ChannelChat::handle_message_sent);
     client.add_model_message_handler(ChannelChat::handle_message_removed);
     client.add_model_message_handler(ChannelChat::handle_message_updated);

crates/client/src/client.rs

@@ -14,22 +14,18 @@ use async_tungstenite::tungstenite::{
 };
 use chrono::{DateTime, Utc};
 use clock::SystemClock;
-use collections::HashMap;
 use futures::{
-    channel::oneshot,
+    channel::oneshot, future::BoxFuture, AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt,
-    future::{BoxFuture, LocalBoxFuture},
+    TryFutureExt as _, TryStreamExt,
-    AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
-};
-use gpui::{
-    actions, AnyModel, AnyWeakModel, AppContext, AsyncAppContext, Global, Model, Task, WeakModel,
 };
+use gpui::{actions, AppContext, AsyncAppContext, Global, Model, Task, WeakModel};
 use http_client::{AsyncBody, HttpClient, HttpClientWithUrl};
 use parking_lot::RwLock;
 use postage::watch;
-use proto::ProtoClient;
+use proto::{AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet};
 use rand::prelude::*;
 use release_channel::{AppVersion, ReleaseChannel};
-use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, PeerId, RequestMessage};
+use rpc::proto::{AnyTypedEnvelope, EnvelopedMessage, PeerId, RequestMessage};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use settings::{Settings, SettingsSources};
@@ -208,6 +204,7 @@ pub struct Client {
     telemetry: Arc<Telemetry>,
     credentials_provider: Arc<dyn CredentialsProvider + Send + Sync + 'static>,
     state: RwLock<ClientState>,
+    handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,
 
     #[allow(clippy::type_complexity)]
     #[cfg(any(test, feature = "test-support"))]
@@ -304,30 +301,7 @@ impl Status {
 struct ClientState {
     credentials: Option<Credentials>,
     status: (watch::Sender<Status>, watch::Receiver<Status>),
-    entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
     _reconnect_task: Option<Task<()>>,
-    entities_by_type_and_remote_id: HashMap<(TypeId, u64), WeakSubscriber>,
-    models_by_message_type: HashMap<TypeId, AnyWeakModel>,
-    entity_types_by_message_type: HashMap<TypeId, TypeId>,
-    #[allow(clippy::type_complexity)]
-    message_handlers: HashMap<
-        TypeId,
-        Arc<
-            dyn Send
-                + Sync
-                + Fn(
-                    AnyModel,
-                    Box<dyn AnyTypedEnvelope>,
-                    &Arc<Client>,
-                    AsyncAppContext,
-                ) -> LocalBoxFuture<'static, Result<()>>,
-        >,
-    >,
-}
-
-enum WeakSubscriber {
-    Entity { handle: AnyWeakModel },
-    Pending(Vec<Box<dyn AnyTypedEnvelope>>),
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -379,12 +353,7 @@ impl Default for ClientState {
         Self {
             credentials: None,
             status: watch::channel_with(Status::SignedOut),
-            entity_id_extractors: Default::default(),
             _reconnect_task: None,
-            models_by_message_type: Default::default(),
-            entities_by_type_and_remote_id: Default::default(),
-            entity_types_by_message_type: Default::default(),
-            message_handlers: Default::default(),
         }
     }
 }
@@ -405,13 +374,13 @@ impl Drop for Subscription {
         match self {
             Subscription::Entity { client, id } => {
                 if let Some(client) = client.upgrade() {
-                    let mut state = client.state.write();
+                    let mut state = client.handler_set.lock();
                     let _ = state.entities_by_type_and_remote_id.remove(id);
                 }
             }
             Subscription::Message { client, id } => {
                 if let Some(client) = client.upgrade() {
-                    let mut state = client.state.write();
+                    let mut state = client.handler_set.lock();
                     let _ = state.entity_types_by_message_type.remove(id);
                     let _ = state.message_handlers.remove(id);
                 }
@@ -430,21 +399,21 @@ pub struct PendingEntitySubscription<T: 'static> {
 impl<T: 'static> PendingEntitySubscription<T> {
     pub fn set_model(mut self, model: &Model<T>, cx: &mut AsyncAppContext) -> Subscription {
         self.consumed = true;
-        let mut state = self.client.state.write();
+        let mut handlers = self.client.handler_set.lock();
         let id = (TypeId::of::<T>(), self.remote_id);
-        let Some(WeakSubscriber::Pending(messages)) =
+        let Some(EntityMessageSubscriber::Pending(messages)) =
-            state.entities_by_type_and_remote_id.remove(&id)
+            handlers.entities_by_type_and_remote_id.remove(&id)
         else {
             unreachable!()
         };
 
-        state.entities_by_type_and_remote_id.insert(
+        handlers.entities_by_type_and_remote_id.insert(
             id,
-            WeakSubscriber::Entity {
+            EntityMessageSubscriber::Entity {
                 handle: model.downgrade().into(),
             },
         );
-        drop(state);
+        drop(handlers);
         for message in messages {
             let client_id = self.client.id();
             let type_name = message.payload_type_name();
@@ -467,8 +436,8 @@ impl<T: 'static> PendingEntitySubscription<T> {
 impl<T: 'static> Drop for PendingEntitySubscription<T> {
     fn drop(&mut self) {
         if !self.consumed {
-            let mut state = self.client.state.write();
+            let mut state = self.client.handler_set.lock();
-            if let Some(WeakSubscriber::Pending(messages)) = state
+            if let Some(EntityMessageSubscriber::Pending(messages)) = state
                 .entities_by_type_and_remote_id
                 .remove(&(TypeId::of::<T>(), self.remote_id))
             {
@@ -549,6 +518,7 @@ impl Client {
             http,
             credentials_provider,
             state: Default::default(),
+            handler_set: Default::default(),
 
             #[cfg(any(test, feature = "test-support"))]
             authenticate: Default::default(),
@@ -592,10 +562,7 @@ impl Client {
     pub fn teardown(&self) {
         let mut state = self.state.write();
         state._reconnect_task.take();
-        state.message_handlers.clear();
+        self.handler_set.lock().clear();
-        state.models_by_message_type.clear();
-        state.entities_by_type_and_remote_id.clear();
-        state.entity_id_extractors.clear();
         self.peer.teardown();
     }
 
@@ -708,14 +675,14 @@ impl Client {
     {
         let id = (TypeId::of::<T>(), remote_id);
 
-        let mut state = self.state.write();
+        let mut state = self.handler_set.lock();
         if state.entities_by_type_and_remote_id.contains_key(&id) {
             return Err(anyhow!("already subscribed to entity"));
         }
 
         state
             .entities_by_type_and_remote_id
-            .insert(id, WeakSubscriber::Pending(Default::default()));
+            .insert(id, EntityMessageSubscriber::Pending(Default::default()));
 
         Ok(PendingEntitySubscription {
             client: self.clone(),
@@ -752,13 +719,13 @@ impl Client {
         E: 'static,
         H: 'static
             + Sync
-            + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F
+            + Fn(Model<E>, TypedEnvelope<M>, AnyProtoClient, AsyncAppContext) -> F
             + Send
             + Sync,
         F: 'static + Future<Output = Result<()>>,
     {
         let message_type_id = TypeId::of::<M>();
-        let mut state = self.state.write();
+        let mut state = self.handler_set.lock();
         state
             .models_by_message_type
             .insert(message_type_id, entity.into());
@@ -803,85 +770,18 @@ impl Client {
         })
     }
 
-    pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
-    where
-        M: EntityMessage,
-        E: 'static,
-        H: 'static + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
-        F: 'static + Future<Output = Result<()>>,
-    {
-        self.add_entity_message_handler::<M, E, _, _>(move |subscriber, message, _, cx| {
-            handler(subscriber.downcast::<E>().unwrap(), message, cx)
-        })
-    }
-
-    fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
-    where
-        M: EntityMessage,
-        E: 'static,
-        H: 'static + Fn(AnyModel, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F + Send + Sync,
-        F: 'static + Future<Output = Result<()>>,
-    {
-        let model_type_id = TypeId::of::<E>();
-        let message_type_id = TypeId::of::<M>();
-
-        let mut state = self.state.write();
-        state
-            .entity_types_by_message_type
-            .insert(message_type_id, model_type_id);
-        state
-            .entity_id_extractors
-            .entry(message_type_id)
-            .or_insert_with(|| {
-                |envelope| {
-                    envelope
-                        .as_any()
-                        .downcast_ref::<TypedEnvelope<M>>()
-                        .unwrap()
-                        .payload
-                        .remote_entity_id()
-                }
-            });
-        let prev_handler = state.message_handlers.insert(
-            message_type_id,
-            Arc::new(move |handle, envelope, client, cx| {
-                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
-                handler(handle, *envelope, client.clone(), cx).boxed_local()
-            }),
-        );
-        if prev_handler.is_some() {
-            panic!("registered handler for the same message twice");
-        }
-    }
-
-    pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
-    where
-        M: EntityMessage + RequestMessage,
-        E: 'static,
-        H: 'static + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
-        F: 'static + Future<Output = Result<M::Response>>,
-    {
-        self.add_entity_message_handler::<M, E, _, _>(move |entity, envelope, client, cx| {
-            Self::respond_to_request::<M, _>(
-                envelope.receipt(),
-                handler(entity.downcast::<E>().unwrap(), envelope, cx),
-                client,
-            )
-        })
-    }
-
     async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
         receipt: Receipt<T>,
         response: F,
-        client: Arc<Self>,
+        client: AnyProtoClient,
     ) -> Result<()> {
         match response.await {
             Ok(response) => {
-                client.respond(receipt, response)?;
+                client.send_response(receipt.message_id, response)?;
                 Ok(())
             }
             Err(error) => {
-                client.respond_with_error(receipt, error.to_proto())?;
+                client.send_response(receipt.message_id, error.to_proto())?;
                 Err(error)
             }
         }
@@ -1541,16 +1441,6 @@ impl Client {
         self.peer.send(self.connection_id()?, message)
     }
 
-    pub fn send_dynamic(
-        &self,
-        envelope: proto::Envelope,
-        message_type: &'static str,
-    ) -> Result<()> {
-        log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
-        let connection_id = self.connection_id()?;
-        self.peer.send_dynamic(connection_id, envelope)
-    }
-
     pub fn request<T: RequestMessage>(
         &self,
         request: T,
@@ -1632,115 +1522,56 @@ impl Client {
         }
     }
 
-    fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
-        log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
-        self.peer.respond(receipt, response)
-    }
-
-    fn respond_with_error<T: RequestMessage>(
-        &self,
-        receipt: Receipt<T>,
-        error: proto::Error,
-    ) -> Result<()> {
-        log::debug!("rpc respond. client_id:{}. name:{}", self.id(), T::NAME);
-        self.peer.respond_with_error(receipt, error)
-    }
-
     fn handle_message(
         self: &Arc<Client>,
         message: Box<dyn AnyTypedEnvelope>,
         cx: &AsyncAppContext,
     ) {
-        let mut state = self.state.write();
+        let sender_id = message.sender_id();
+        let request_id = message.message_id();
         let type_name = message.payload_type_name();
-        let payload_type_id = message.payload_type_id();
+        let original_sender_id = message.original_sender_id();
-        let sender_id = message.original_sender_id();
 
-        let mut subscriber = None;
+        if let Some(future) = ProtoMessageHandlerSet::handle_message(
-
+            &self.handler_set,
-        if let Some(handle) = state
+            message,
-            .models_by_message_type
+            self.clone().into(),
-            .get(&payload_type_id)
+            cx.clone(),
-            .and_then(|handle| handle.upgrade())
+        ) {
-        {
-            subscriber = Some(handle);
-        } else if let Some((extract_entity_id, entity_type_id)) =
-            state.entity_id_extractors.get(&payload_type_id).zip(
-                state
-                    .entity_types_by_message_type
-                    .get(&payload_type_id)
-                    .copied(),
-            )
-        {
-            let entity_id = (extract_entity_id)(message.as_ref());
-
-            match state
-                .entities_by_type_and_remote_id
-                .get_mut(&(entity_type_id, entity_id))
-            {
-                Some(WeakSubscriber::Pending(pending)) => {
-                    pending.push(message);
-                    return;
-                }
-                Some(weak_subscriber) => match weak_subscriber {
-                    WeakSubscriber::Entity { handle } => {
-                        subscriber = handle.upgrade();
-                    }
-
-                    WeakSubscriber::Pending(_) => {}
-                },
-                _ => {}
-            }
-        }
-
-        let subscriber = if let Some(subscriber) = subscriber {
-            subscriber
-        } else {
-            log::info!("unhandled message {}", type_name);
-            self.peer.respond_with_unhandled_message(message).log_err();
-            return;
-        };
-
-        let handler = state.message_handlers.get(&payload_type_id).cloned();
-        // Dropping the state prevents deadlocks if the handler interacts with rpc::Client.
-        // It also ensures we don't hold the lock while yielding back to the executor, as
-        // that might cause the executor thread driving this future to block indefinitely.
-        drop(state);
-
-        if let Some(handler) = handler {
-            let future = handler(subscriber, message, self, cx.clone());
             let client_id = self.id();
             log::debug!(
                 "rpc message received. client_id:{}, sender_id:{:?}, type:{}",
                 client_id,
-                sender_id,
+                original_sender_id,
                 type_name
             );
             cx.spawn(move |_| async move {
-                    match future.await {
+                match future.await {
-                        Ok(()) => {
+                    Ok(()) => {
-                            log::debug!(
+                        log::debug!(
-                                "rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
+                            "rpc message handled. client_id:{}, sender_id:{:?}, type:{}",
-                                client_id,
+                            client_id,
-                                sender_id,
+                            original_sender_id,
-                                type_name
+                            type_name
-                            );
+                        );
-                        }
-                        Err(error) => {
-                            log::error!(
-                                "error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
-                                client_id,
-                                sender_id,
-                                type_name,
-                                error
-                            );
-                        }
                     }
-                })
+                    Err(error) => {
-                .detach();
+                        log::error!(
+                            "error handling message. client_id:{}, sender_id:{:?}, type:{}, error:{:?}",
+                            client_id,
+                            original_sender_id,
+                            type_name,
+                            error
+                        );
+                    }
+                }
+            })
+            .detach();
         } else {
             log::info!("unhandled message {}", type_name);
-            self.peer.respond_with_unhandled_message(message).log_err();
+            self.peer
+                .respond_with_unhandled_message(sender_id.into(), request_id, type_name)
+                .log_err();
         }
     }
 
@@ -1759,7 +1590,23 @@ impl ProtoClient for Client {
     }
 
     fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
-        self.send_dynamic(envelope, message_type)
+        log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
+        let connection_id = self.connection_id()?;
+        self.peer.send_dynamic(connection_id, envelope)
+    }
+
+    fn send_response(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
+        log::debug!(
+            "rpc respond. client_id:{}, name:{}",
+            self.id(),
+            message_type
+        );
+        let connection_id = self.connection_id()?;
+        self.peer.send_dynamic(connection_id, envelope)
+    }
+
+    fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet> {
+        &self.handler_set
     }
 }
 
@@ -2103,7 +1950,7 @@ mod tests {
 
         let (done_tx1, mut done_rx1) = smol::channel::unbounded();
         let (done_tx2, mut done_rx2) = smol::channel::unbounded();
-        client.add_model_message_handler(
+        AnyProtoClient::from(client.clone()).add_model_message_handler(
             move |model: Model<TestModel>, _: TypedEnvelope<proto::JoinProject>, mut cx| {
                 match model.update(&mut cx, |model, _| model.id).unwrap() {
                     1 => done_tx1.try_send(()).unwrap(),

crates/collab/src/tests/test_server.rs

@@ -301,7 +301,7 @@ impl TestServer {
             dev_server_projects::init(client.clone(), cx);
             settings::KeymapFile::load_asset(os_keymap, cx).unwrap();
             language_model::LanguageModelRegistry::test(cx);
-            assistant::context_store::init(&client);
+            assistant::context_store::init(&client.clone().into());
         });
 
         client

crates/project/src/buffer_store.rs

@@ -71,7 +71,7 @@ pub struct ProjectTransaction(pub HashMap<Model<Buffer>, language::Transaction>)
 impl EventEmitter<BufferStoreEvent> for BufferStore {}
 
 impl BufferStore {
-    pub fn init(client: &Arc<Client>) {
+    pub fn init(client: &AnyProtoClient) {
         client.add_model_message_handler(Self::handle_buffer_reloaded);
         client.add_model_message_handler(Self::handle_buffer_saved);
         client.add_model_message_handler(Self::handle_update_buffer_file);

crates/project/src/project.rs

@@ -65,7 +65,10 @@ use paths::{
 use prettier_support::{DefaultPrettier, PrettierInstance};
 use project_settings::{LspSettings, ProjectSettings};
 use remote::SshSession;
-use rpc::{proto::AnyProtoClient, ErrorCode};
+use rpc::{
+    proto::{AnyProtoClient, SSH_PROJECT_ID},
+    ErrorCode,
+};
 use search::{SearchQuery, SearchResult};
 use search_history::SearchHistory;
 use settings::{watch_config_file, Settings, SettingsLocation, SettingsStore};
@@ -574,6 +577,7 @@ impl Project {
         connection_manager::init(client.clone(), cx);
         Self::init_settings(cx);
 
+        let client: AnyProtoClient = client.clone().into();
         client.add_model_message_handler(Self::handle_add_collaborator);
         client.add_model_message_handler(Self::handle_update_project_collaborator);
         client.add_model_message_handler(Self::handle_remove_collaborator);
@@ -594,9 +598,9 @@ impl Project {
         client.add_model_request_handler(Self::handle_task_templates);
         client.add_model_message_handler(Self::handle_create_buffer_for_peer);
 
-        WorktreeStore::init(client);
+        WorktreeStore::init(&client);
-        BufferStore::init(client);
+        BufferStore::init(&client);
-        LspStore::init(client);
+        LspStore::init(&client);
     }
 
     pub fn local(
@@ -697,15 +701,19 @@ impl Project {
     ) -> Model<Self> {
         let this = Self::local(client, node, user_store, languages, fs, None, cx);
         this.update(cx, |this, cx| {
-            let buffer_store = this.buffer_store.downgrade();
+            let client: AnyProtoClient = ssh.clone().into();
+
             this.worktree_store.update(cx, |store, _cx| {
-                store.set_upstream_client(ssh.clone().into());
+                store.set_upstream_client(client.clone());
             });
 
-            ssh.add_message_handler(cx.weak_model(), Self::handle_update_worktree);
+            ssh.subscribe_to_entity(SSH_PROJECT_ID, &cx.handle());
-            ssh.add_message_handler(cx.weak_model(), Self::handle_create_buffer_for_peer);
+            ssh.subscribe_to_entity(SSH_PROJECT_ID, &this.buffer_store);
-            ssh.add_message_handler(buffer_store.clone(), BufferStore::handle_update_buffer_file);
+            ssh.subscribe_to_entity(SSH_PROJECT_ID, &this.worktree_store);
-            ssh.add_message_handler(buffer_store.clone(), BufferStore::handle_update_diff_base);
+            client.add_model_message_handler(Self::handle_update_worktree);
+            client.add_model_message_handler(Self::handle_create_buffer_for_peer);
+            client.add_model_message_handler(BufferStore::handle_update_buffer_file);
+            client.add_model_message_handler(BufferStore::handle_update_diff_base);
 
             this.ssh_session = Some(ssh);
         });

crates/project/src/worktree_store.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use anyhow::{anyhow, Context as _, Result};
-use client::{Client, DevServerProjectId};
+use client::DevServerProjectId;
 use collections::{HashMap, HashSet};
 use fs::Fs;
 use futures::{
@@ -17,7 +17,7 @@ use gpui::{
 };
 use postage::oneshot;
 use rpc::{
-    proto::{self, AnyProtoClient},
+    proto::{self, AnyProtoClient, SSH_PROJECT_ID},
     TypedEnvelope,
 };
 use smol::{
@@ -58,12 +58,12 @@ pub enum WorktreeStoreEvent {
 impl EventEmitter<WorktreeStoreEvent> for WorktreeStore {}
 
 impl WorktreeStore {
-    pub fn init(client: &Arc<Client>) {
+    pub fn init(client: &AnyProtoClient) {
-        client.add_model_request_handler(WorktreeStore::handle_create_project_entry);
+        client.add_model_request_handler(Self::handle_create_project_entry);
-        client.add_model_request_handler(WorktreeStore::handle_rename_project_entry);
+        client.add_model_request_handler(Self::handle_rename_project_entry);
-        client.add_model_request_handler(WorktreeStore::handle_copy_project_entry);
+        client.add_model_request_handler(Self::handle_copy_project_entry);
-        client.add_model_request_handler(WorktreeStore::handle_delete_project_entry);
+        client.add_model_request_handler(Self::handle_delete_project_entry);
-        client.add_model_request_handler(WorktreeStore::handle_expand_project_entry);
+        client.add_model_request_handler(Self::handle_expand_project_entry);
     }
 
     pub fn new(retain_worktrees: bool, fs: Arc<dyn Fs>) -> Self {
@@ -188,7 +188,10 @@ impl WorktreeStore {
         let path = abs_path.to_string_lossy().to_string();
         cx.spawn(|this, mut cx| async move {
             let response = client
-                .request(proto::AddWorktree { path: path.clone() })
+                .request(proto::AddWorktree {
+                    project_id: SSH_PROJECT_ID,
+                    path: path.clone(),
+                })
                 .await?;
             let worktree = cx.update(|cx| {
                 Worktree::remote(

crates/proto/Cargo.toml

@@ -20,8 +20,10 @@ doctest = false
 anyhow.workspace = true
 collections.workspace = true
 futures.workspace = true
+parking_lot.workspace = true
 prost.workspace = true
 serde.workspace = true
+gpui.workspace = true
 
 [build-dependencies]
 prost-build.workspace = true

crates/proto/proto/zed.proto

@@ -2484,6 +2484,7 @@ message GetLlmTokenResponse {
 // Remote FS
 
 message AddWorktree {
+    uint64 project_id = 2;
     string path = 1;
 }
 

crates/proto/src/proto.rs

@@ -2,14 +2,14 @@
 
 pub mod error;
 mod macros;
+mod proto_client;
 mod typed_envelope;
 
 pub use error::*;
+pub use proto_client::*;
 pub use typed_envelope::*;
 
-use anyhow::anyhow;
 use collections::HashMap;
-use futures::{future::BoxFuture, Future};
 pub use prost::{DecodeError, Message};
 use serde::Serialize;
 use std::{
@@ -17,12 +17,14 @@ use std::{
     cmp,
     fmt::{self, Debug},
     iter, mem,
-    sync::Arc,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 include!(concat!(env!("OUT_DIR"), "/zed.messages.rs"));
 
+pub const SSH_PEER_ID: PeerId = PeerId { owner_id: 0, id: 0 };
+pub const SSH_PROJECT_ID: u64 = 0;
+
 pub trait EnvelopedMessage: Clone + Debug + Serialize + Sized + Send + Sync + 'static {
     const NAME: &'static str;
     const PRIORITY: MessagePriority;
@@ -60,51 +62,6 @@ pub enum MessagePriority {
     Background,
 }
 
-pub trait ProtoClient: Send + Sync {
-    fn request(
-        &self,
-        envelope: Envelope,
-        request_type: &'static str,
-    ) -> BoxFuture<'static, anyhow::Result<Envelope>>;
-
-    fn send(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
-}
-
-#[derive(Clone)]
-pub struct AnyProtoClient(Arc<dyn ProtoClient>);
-
-impl<T> From<Arc<T>> for AnyProtoClient
-where
-    T: ProtoClient + 'static,
-{
-    fn from(client: Arc<T>) -> Self {
-        Self(client)
-    }
-}
-
-impl AnyProtoClient {
-    pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
-        Self(client)
-    }
-
-    pub fn request<T: RequestMessage>(
-        &self,
-        request: T,
-    ) -> impl Future<Output = anyhow::Result<T::Response>> {
-        let envelope = request.into_envelope(0, None, None);
-        let response = self.0.request(envelope, T::NAME);
-        async move {
-            T::Response::from_envelope(response.await?)
-                .ok_or_else(|| anyhow!("received response of the wrong type"))
-        }
-    }
-
-    pub fn send<T: EnvelopedMessage>(&self, request: T) -> anyhow::Result<()> {
-        let envelope = request.into_envelope(0, None, None);
-        self.0.send(envelope, T::NAME)
-    }
-}
-
 impl<T: EnvelopedMessage> AnyTypedEnvelope for TypedEnvelope<T> {
     fn payload_type_id(&self) -> TypeId {
         TypeId::of::<T>()
@@ -537,11 +494,13 @@ request_messages!(
 entity_messages!(
     {project_id, ShareProject},
     AddProjectCollaborator,
+    AddWorktree,
     ApplyCodeAction,
     ApplyCompletionAdditionalEdits,
     BlameBuffer,
     BufferReloaded,
     BufferSaved,
+    CloseBuffer,
     CopyProjectEntry,
     CreateBufferForPeer,
     CreateProjectEntry,

crates/proto/src/proto_client.rs

@@ -0,0 +1,277 @@
+use crate::{
+    error::ErrorExt as _, AnyTypedEnvelope, EntityMessage, Envelope, EnvelopedMessage,
+    RequestMessage, TypedEnvelope,
+};
+use anyhow::anyhow;
+use collections::HashMap;
+use futures::{
+    future::{BoxFuture, LocalBoxFuture},
+    Future, FutureExt as _,
+};
+use gpui::{AnyModel, AnyWeakModel, AsyncAppContext, Model};
+pub use prost::Message;
+use std::{any::TypeId, sync::Arc};
+
+#[derive(Clone)]
+pub struct AnyProtoClient(Arc<dyn ProtoClient>);
+
+pub trait ProtoClient: Send + Sync {
+    fn request(
+        &self,
+        envelope: Envelope,
+        request_type: &'static str,
+    ) -> BoxFuture<'static, anyhow::Result<Envelope>>;
+
+    fn send(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
+
+    fn send_response(&self, envelope: Envelope, message_type: &'static str) -> anyhow::Result<()>;
+
+    fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet>;
+}
+
+#[derive(Default)]
+pub struct ProtoMessageHandlerSet {
+    pub entity_types_by_message_type: HashMap<TypeId, TypeId>,
+    pub entities_by_type_and_remote_id: HashMap<(TypeId, u64), EntityMessageSubscriber>,
+    pub entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
+    pub models_by_message_type: HashMap<TypeId, AnyWeakModel>,
+    pub message_handlers: HashMap<TypeId, ProtoMessageHandler>,
+}
+
+pub type ProtoMessageHandler = Arc<
+    dyn Send
+        + Sync
+        + Fn(
+            AnyModel,
+            Box<dyn AnyTypedEnvelope>,
+            AnyProtoClient,
+            AsyncAppContext,
+        ) -> LocalBoxFuture<'static, anyhow::Result<()>>,
+>;
+
+impl ProtoMessageHandlerSet {
+    pub fn clear(&mut self) {
+        self.message_handlers.clear();
+        self.models_by_message_type.clear();
+        self.entities_by_type_and_remote_id.clear();
+        self.entity_id_extractors.clear();
+    }
+
+    fn add_message_handler(
+        &mut self,
+        message_type_id: TypeId,
+        model: gpui::AnyWeakModel,
+        handler: ProtoMessageHandler,
+    ) {
+        self.models_by_message_type.insert(message_type_id, model);
+        let prev_handler = self.message_handlers.insert(message_type_id, handler);
+        if prev_handler.is_some() {
+            panic!("registered handler for the same message twice");
+        }
+    }
+
+    fn add_entity_message_handler(
+        &mut self,
+        message_type_id: TypeId,
+        model_type_id: TypeId,
+        entity_id_extractor: fn(&dyn AnyTypedEnvelope) -> u64,
+        handler: ProtoMessageHandler,
+    ) {
+        self.entity_id_extractors
+            .entry(message_type_id)
+            .or_insert(entity_id_extractor);
+        self.entity_types_by_message_type
+            .insert(message_type_id, model_type_id);
+        let prev_handler = self.message_handlers.insert(message_type_id, handler);
+        if prev_handler.is_some() {
+            panic!("registered handler for the same message twice");
+        }
+    }
+
+    pub fn handle_message(
+        this: &parking_lot::Mutex<Self>,
+        message: Box<dyn AnyTypedEnvelope>,
+        client: AnyProtoClient,
+        cx: AsyncAppContext,
+    ) -> Option<LocalBoxFuture<'static, anyhow::Result<()>>> {
+        let payload_type_id = message.payload_type_id();
+        let mut this = this.lock();
+        let handler = this.message_handlers.get(&payload_type_id)?.clone();
+        let entity = if let Some(entity) = this.models_by_message_type.get(&payload_type_id) {
+            entity.upgrade()?
+        } else {
+            let extract_entity_id = *this.entity_id_extractors.get(&payload_type_id)?;
+            let entity_type_id = *this.entity_types_by_message_type.get(&payload_type_id)?;
+            let entity_id = (extract_entity_id)(message.as_ref());
+
+            match this
+                .entities_by_type_and_remote_id
+                .get_mut(&(entity_type_id, entity_id))?
+            {
+                EntityMessageSubscriber::Pending(pending) => {
+                    pending.push(message);
+                    return None;
+                }
+                EntityMessageSubscriber::Entity { handle } => handle.upgrade()?,
+            }
+        };
+        drop(this);
+        Some(handler(entity, message, client, cx))
+    }
+}
+
+pub enum EntityMessageSubscriber {
+    Entity { handle: AnyWeakModel },
+    Pending(Vec<Box<dyn AnyTypedEnvelope>>),
+}
+
+impl<T> From<Arc<T>> for AnyProtoClient
+where
+    T: ProtoClient + 'static,
+{
+    fn from(client: Arc<T>) -> Self {
+        Self(client)
+    }
+}
+
+impl AnyProtoClient {
+    pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
+        Self(client)
+    }
+
+    pub fn request<T: RequestMessage>(
+        &self,
+        request: T,
+    ) -> impl Future<Output = anyhow::Result<T::Response>> {
+        let envelope = request.into_envelope(0, None, None);
+        let response = self.0.request(envelope, T::NAME);
+        async move {
+            T::Response::from_envelope(response.await?)
+                .ok_or_else(|| anyhow!("received response of the wrong type"))
+        }
+    }
+
+    pub fn send<T: EnvelopedMessage>(&self, request: T) -> anyhow::Result<()> {
+        let envelope = request.into_envelope(0, None, None);
+        self.0.send(envelope, T::NAME)
+    }
+
+    pub fn send_response<T: EnvelopedMessage>(
+        &self,
+        request_id: u32,
+        request: T,
+    ) -> anyhow::Result<()> {
+        let envelope = request.into_envelope(0, Some(request_id), None);
+        self.0.send(envelope, T::NAME)
+    }
+
+    pub fn add_request_handler<M, E, H, F>(&self, model: gpui::WeakModel<E>, handler: H)
+    where
+        M: RequestMessage,
+        E: 'static,
+        H: 'static + Sync + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F + Send + Sync,
+        F: 'static + Future<Output = anyhow::Result<M::Response>>,
+    {
+        self.0.message_handler_set().lock().add_message_handler(
+            TypeId::of::<M>(),
+            model.into(),
+            Arc::new(move |model, envelope, client, cx| {
+                let model = model.downcast::<E>().unwrap();
+                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
+                let request_id = envelope.message_id();
+                handler(model, *envelope, cx)
+                    .then(move |result| async move {
+                        match result {
+                            Ok(response) => {
+                                client.send_response(request_id, response)?;
+                                Ok(())
+                            }
+                            Err(error) => {
+                                client.send_response(request_id, error.to_proto())?;
+                                Err(error)
+                            }
+                        }
+                    })
+                    .boxed_local()
+            }),
+        )
+    }
+
+    pub fn add_model_request_handler<M, E, H, F>(&self, handler: H)
+    where
+        M: EnvelopedMessage + RequestMessage + EntityMessage,
+        E: 'static,
+        H: 'static + Sync + Send + Fn(gpui::Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
+        F: 'static + Future<Output = anyhow::Result<M::Response>>,
+    {
+        let message_type_id = TypeId::of::<M>();
+        let model_type_id = TypeId::of::<E>();
+        let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
+            envelope
+                .as_any()
+                .downcast_ref::<TypedEnvelope<M>>()
+                .unwrap()
+                .payload
+                .remote_entity_id()
+        };
+        self.0
+            .message_handler_set()
+            .lock()
+            .add_entity_message_handler(
+                message_type_id,
+                model_type_id,
+                entity_id_extractor,
+                Arc::new(move |model, envelope, client, cx| {
+                    let model = model.downcast::<E>().unwrap();
+                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
+                    let request_id = envelope.message_id();
+                    handler(model, *envelope, cx)
+                        .then(move |result| async move {
+                            match result {
+                                Ok(response) => {
+                                    client.send_response(request_id, response)?;
+                                    Ok(())
+                                }
+                                Err(error) => {
+                                    client.send_response(request_id, error.to_proto())?;
+                                    Err(error)
+                                }
+                            }
+                        })
+                        .boxed_local()
+                }),
+            );
+    }
+
+    pub fn add_model_message_handler<M, E, H, F>(&self, handler: H)
+    where
+        M: EnvelopedMessage + EntityMessage,
+        E: 'static,
+        H: 'static + Sync + Send + Fn(gpui::Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
+        F: 'static + Future<Output = anyhow::Result<()>>,
+    {
+        let message_type_id = TypeId::of::<M>();
+        let model_type_id = TypeId::of::<E>();
+        let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
+            envelope
+                .as_any()
+                .downcast_ref::<TypedEnvelope<M>>()
+                .unwrap()
+                .payload
+                .remote_entity_id()
+        };
+        self.0
+            .message_handler_set()
+            .lock()
+            .add_entity_message_handler(
+                message_type_id,
+                model_type_id,
+                entity_id_extractor,
+                Arc::new(move |model, envelope, _, cx| {
+                    let model = model.downcast::<E>().unwrap();
+                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
+                    handler(model, *envelope, cx).boxed_local()
+                }),
+            );
+    }
+}

crates/remote/src/ssh_session.rs

@@ -8,17 +8,14 @@ use anyhow::{anyhow, Context as _, Result};
 use collections::HashMap;
 use futures::{
     channel::{mpsc, oneshot},
-    future::{BoxFuture, LocalBoxFuture},
+    future::BoxFuture,
     select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, StreamExt as _,
 };
-use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion, WeakModel};
+use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion};
 use parking_lot::Mutex;
-use rpc::{
+use rpc::proto::{
-    proto::{
+    self, build_typed_envelope, EntityMessageSubscriber, Envelope, EnvelopedMessage, PeerId,
-        self, build_typed_envelope, AnyTypedEnvelope, Envelope, EnvelopedMessage, PeerId,
+    ProtoClient, ProtoMessageHandlerSet, RequestMessage,
-        ProtoClient, RequestMessage,
-    },
-    TypedEnvelope,
 };
 use smol::{
     fs,
@@ -48,20 +45,7 @@ pub struct SshSession {
     outgoing_tx: mpsc::UnboundedSender<Envelope>,
     spawn_process_tx: mpsc::UnboundedSender<SpawnRequest>,
     client_socket: Option<SshSocket>,
-    message_handlers: Mutex<
+    state: Mutex<ProtoMessageHandlerSet>,
-        HashMap<
-            TypeId,
-            Arc<
-                dyn Send
-                    + Sync
-                    + Fn(
-                        Box<dyn AnyTypedEnvelope>,
-                        Arc<SshSession>,
-                        AsyncAppContext,
-                    ) -> Option<LocalBoxFuture<'static, Result<()>>>,
-            >,
-        >,
-    >,
 }
 
 struct SshClientState {
@@ -330,7 +314,7 @@ impl SshSession {
             outgoing_tx,
             spawn_process_tx,
             client_socket,
-            message_handlers: Default::default(),
+            state: Default::default(),
         });
 
         cx.spawn(|cx| {
@@ -351,18 +335,26 @@ impl SshSession {
                     } else if let Some(envelope) =
                         build_typed_envelope(peer_id, Instant::now(), incoming)
                     {
-                        log::debug!(
+                        let type_name = envelope.payload_type_name();
-                            "ssh message received. name:{}",
+                        if let Some(future) = ProtoMessageHandlerSet::handle_message(
-                            envelope.payload_type_name()
+                            &this.state,
-                        );
+                            envelope,
-                        let type_id = envelope.payload_type_id();
+                            this.clone().into(),
-                        let handler = this.message_handlers.lock().get(&type_id).cloned();
+                            cx.clone(),
-                        if let Some(handler) = handler {
+                        ) {
-                            if let Some(future) = handler(envelope, this.clone(), cx.clone()) {
+                            log::debug!("ssh message received. name:{type_name}");
-                                future.await.ok();
+                            match future.await {
-                            } else {
+                                Ok(_) => {
-                                this.message_handlers.lock().remove(&type_id);
+                                    log::debug!("ssh message handled. name:{type_name}");
+                                }
+                                Err(error) => {
+                                    log::error!(
+                                        "error handling message. type:{type_name}, error:{error:?}",
+                                    );
+                                }
                             }
+                        } else {
+                            log::error!("unhandled ssh message name:{type_name}");
                         }
                     }
                 }
@@ -389,6 +381,7 @@ impl SshSession {
     }
 
     pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
+        log::debug!("ssh send name:{}", T::NAME);
         self.send_dynamic(payload.into_envelope(0, None, None))
     }
 
@@ -412,6 +405,22 @@ impl SshSession {
         Ok(())
     }
 
+    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
+        let id = (TypeId::of::<E>(), remote_id);
+
+        let mut state = self.state.lock();
+        if state.entities_by_type_and_remote_id.contains_key(&id) {
+            panic!("already subscribed to entity");
+        }
+
+        state.entities_by_type_and_remote_id.insert(
+            id,
+            EntityMessageSubscriber::Entity {
+                handle: entity.downgrade().into(),
+            },
+        );
+    }
+
     pub async fn spawn_process(&self, command: String) -> process::Child {
         let (process_tx, process_rx) = oneshot::channel();
         self.spawn_process_tx
@@ -426,54 +435,6 @@ impl SshSession {
     pub fn ssh_args(&self) -> Vec<String> {
         self.client_socket.as_ref().unwrap().ssh_args()
     }
-
-    pub fn add_message_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)
-    where
-        M: EnvelopedMessage,
-        E: 'static,
-        H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
-        F: 'static + Future<Output = Result<()>>,
-    {
-        let message_type_id = TypeId::of::<M>();
-        self.message_handlers.lock().insert(
-            message_type_id,
-            Arc::new(move |envelope, _, cx| {
-                let entity = entity.upgrade()?;
-                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
-                Some(handler(entity, *envelope, cx).boxed_local())
-            }),
-        );
-    }
-
-    pub fn add_request_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)
-    where
-        M: EnvelopedMessage + RequestMessage,
-        E: 'static,
-        H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,
-        F: 'static + Future<Output = Result<M::Response>>,
-    {
-        let message_type_id = TypeId::of::<M>();
-        self.message_handlers.lock().insert(
-            message_type_id,
-            Arc::new(move |envelope, this, cx| {
-                let entity = entity.upgrade()?;
-                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
-                let request_id = envelope.message_id();
-                Some(
-                    handler(entity, *envelope, cx)
-                        .then(move |result| async move {
-                            this.outgoing_tx.unbounded_send(result?.into_envelope(
-                                this.next_message_id.fetch_add(1, SeqCst),
-                                Some(request_id),
-                                None,
-                            ))?;
-                            Ok(())
-                        })
-                        .boxed_local(),
-                )
-            }),
-        );
-    }
 }
 
 impl ProtoClient for SshSession {
@@ -488,6 +449,14 @@ impl ProtoClient for SshSession {
     fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
         self.send_dynamic(envelope)
     }
+
+    fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
+        self.send_dynamic(envelope)
+    }
+
+    fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
+        &self.state
+    }
 }
 
 impl SshClientState {

crates/remote_server/src/headless_project.rs

@@ -7,7 +7,7 @@ use project::{
 };
 use remote::SshSession;
 use rpc::{
-    proto::{self, AnyProtoClient, PeerId},
+    proto::{self, AnyProtoClient, SSH_PEER_ID, SSH_PROJECT_ID},
     TypedEnvelope,
 };
 use settings::{Settings as _, SettingsStore};
@@ -18,9 +18,6 @@ use std::{
 };
 use worktree::Worktree;
 
-const PEER_ID: PeerId = PeerId { owner_id: 0, id: 0 };
-const PROJECT_ID: u64 = 0;
-
 pub struct HeadlessProject {
     pub fs: Arc<dyn Fs>,
     pub session: AnyProtoClient,
@@ -36,48 +33,34 @@ impl HeadlessProject {
     }
 
     pub fn new(session: Arc<SshSession>, fs: Arc<dyn Fs>, cx: &mut ModelContext<Self>) -> Self {
-        let this = cx.weak_model();
-
         let worktree_store = cx.new_model(|_| WorktreeStore::new(true, fs.clone()));
         let buffer_store = cx.new_model(|cx| {
-            let mut buffer_store = BufferStore::new(worktree_store.clone(), Some(PROJECT_ID), cx);
+            let mut buffer_store =
-            buffer_store.shared(PROJECT_ID, session.clone().into(), cx);
+                BufferStore::new(worktree_store.clone(), Some(SSH_PROJECT_ID), cx);
+            buffer_store.shared(SSH_PROJECT_ID, session.clone().into(), cx);
             buffer_store
         });
 
-        session.add_request_handler(this.clone(), Self::handle_list_remote_directory);
+        let client: AnyProtoClient = session.clone().into();
-        session.add_request_handler(this.clone(), Self::handle_add_worktree);
-        session.add_request_handler(this.clone(), Self::handle_open_buffer_by_path);
-        session.add_request_handler(this.clone(), Self::handle_find_search_candidates);
 
-        session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_blame_buffer);
+        session.subscribe_to_entity(SSH_PROJECT_ID, &worktree_store);
-        session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_update_buffer);
+        session.subscribe_to_entity(SSH_PROJECT_ID, &buffer_store);
-        session.add_request_handler(buffer_store.downgrade(), BufferStore::handle_save_buffer);
+        session.subscribe_to_entity(SSH_PROJECT_ID, &cx.handle());
-        session.add_message_handler(buffer_store.downgrade(), BufferStore::handle_close_buffer);
 
-        session.add_request_handler(
+        client.add_request_handler(cx.weak_model(), Self::handle_list_remote_directory);
-            worktree_store.downgrade(),
+
-            WorktreeStore::handle_create_project_entry,
+        client.add_model_request_handler(Self::handle_add_worktree);
-        );
+        client.add_model_request_handler(Self::handle_open_buffer_by_path);
-        session.add_request_handler(
+        client.add_model_request_handler(Self::handle_find_search_candidates);
-            worktree_store.downgrade(),
+
-            WorktreeStore::handle_rename_project_entry,
+        client.add_model_request_handler(BufferStore::handle_update_buffer);
-        );
+        client.add_model_message_handler(BufferStore::handle_close_buffer);
-        session.add_request_handler(
+
-            worktree_store.downgrade(),
+        BufferStore::init(&client);
-            WorktreeStore::handle_copy_project_entry,
+        WorktreeStore::init(&client);
-        );
-        session.add_request_handler(
-            worktree_store.downgrade(),
-            WorktreeStore::handle_delete_project_entry,
-        );
-        session.add_request_handler(
-            worktree_store.downgrade(),
-            WorktreeStore::handle_expand_project_entry,
-        );
 
         HeadlessProject {
-            session: session.into(),
+            session: client,
             fs,
             worktree_store,
             buffer_store,
@@ -144,7 +127,7 @@ impl HeadlessProject {
         let buffer_id = buffer.read_with(&cx, |b, _| b.remote_id())?;
         buffer_store.update(&mut cx, |buffer_store, cx| {
             buffer_store
-                .create_buffer_for_peer(&buffer, PEER_ID, cx)
+                .create_buffer_for_peer(&buffer, SSH_PEER_ID, cx)
                 .detach_and_log_err(cx);
         })?;
 
@@ -181,7 +164,7 @@ impl HeadlessProject {
             response.buffer_ids.push(buffer_id.to_proto());
             buffer_store
                 .update(&mut cx, |buffer_store, cx| {
-                    buffer_store.create_buffer_for_peer(&buffer, PEER_ID, cx)
+                    buffer_store.create_buffer_for_peer(&buffer, SSH_PEER_ID, cx)
                 })?
                 .await?;
         }

crates/remote_server/src/remote_editing_tests.rs

@@ -17,7 +17,7 @@ use smol::stream::StreamExt;
 use std::{path::Path, sync::Arc};
 
 #[gpui::test]
-async fn test_remote_editing(cx: &mut TestAppContext, server_cx: &mut TestAppContext) {
+async fn test_basic_remote_editing(cx: &mut TestAppContext, server_cx: &mut TestAppContext) {
     let (project, _headless, fs) = init_test(cx, server_cx).await;
     let (worktree, _) = project
         .update(cx, |project, cx| {

crates/rpc/src/peer.rs

@@ -639,14 +639,13 @@ impl Peer {
 
     pub fn respond_with_unhandled_message(
         &self,
-        envelope: Box<dyn AnyTypedEnvelope>,
+        sender_id: ConnectionId,
+        request_message_id: u32,
+        message_type_name: &'static str,
     ) -> Result<()> {
-        let connection = self.connection_state(envelope.sender_id().into())?;
+        let connection = self.connection_state(sender_id)?;
         let response = ErrorCode::Internal
-            .message(format!(
+            .message(format!("message {} was not handled", message_type_name))
-                "message {} was not handled",
-                envelope.payload_type_name()
-            ))
             .to_proto();
         let message_id = connection
             .next_message_id
@@ -655,7 +654,7 @@ impl Peer {
             .outgoing_tx
             .unbounded_send(proto::Message::Envelope(response.into_envelope(
                 message_id,
-                Some(envelope.message_id()),
+                Some(request_message_id),
                 None,
             )))?;
         Ok(())