ZIm/crates/remote_server/src/headless_project.rs

use anyhow::Result;
use fs::Fs;
use gpui::{AppContext, AsyncAppContext, Context, Model, ModelContext};
use project::{
    buffer_store::{BufferStore, BufferStoreEvent},
    worktree_store::WorktreeStore,
    ProjectPath, WorktreeId, WorktreeSettings,
};
use remote::SshSession;
use rpc::{
    proto::{self, AnyProtoClient, PeerId},
    TypedEnvelope,
};
use settings::{Settings as _, SettingsStore};
use std::{
    path::{Path, PathBuf},
    sync::{atomic::AtomicUsize, Arc},
};
use util::ResultExt as _;
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,
    pub worktree_store: Model<WorktreeStore>,
    pub buffer_store: Model<BufferStore>,
    pub next_entry_id: Arc<AtomicUsize>,
}

impl HeadlessProject {
    pub fn init(cx: &mut AppContext) {
        cx.set_global(SettingsStore::new(cx));
        WorktreeSettings::register(cx);
    }

    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));
        let buffer_store = cx.new_model(|cx| BufferStore::new(worktree_store.clone(), true, cx));
        cx.subscribe(&buffer_store, Self::on_buffer_store_event)
            .detach();

        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_update_buffer);
        session.add_request_handler(this.clone(), Self::handle_save_buffer);
        session.add_request_handler(
            worktree_store.downgrade(),
            WorktreeStore::handle_create_project_entry,
        );
        session.add_request_handler(
            worktree_store.downgrade(),
            WorktreeStore::handle_rename_project_entry,
        );
        session.add_request_handler(
            worktree_store.downgrade(),
            WorktreeStore::handle_copy_project_entry,
        );
        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(),
            fs,
            worktree_store,
            buffer_store,
            next_entry_id: Default::default(),
        }
    }

    pub async fn handle_add_worktree(
        this: Model<Self>,
        message: TypedEnvelope<proto::AddWorktree>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::AddWorktreeResponse> {
        let worktree = this
            .update(&mut cx.clone(), |this, _| {
                Worktree::local(
                    Path::new(&message.payload.path),
                    true,
                    this.fs.clone(),
                    this.next_entry_id.clone(),
                    &mut cx,
                )
            })?
            .await?;

        this.update(&mut cx, |this, cx| {
            let session = this.session.clone();
            this.worktree_store.update(cx, |worktree_store, cx| {
                worktree_store.add(&worktree, cx);
            });
            worktree.update(cx, |worktree, cx| {
                worktree.observe_updates(0, cx, move |update| {
                    session.send(update).ok();
                    futures::future::ready(true)
                });
                proto::AddWorktreeResponse {
                    worktree_id: worktree.id().to_proto(),
                }
            })
        })
    }

    pub async fn handle_update_buffer(
        this: Model<Self>,
        envelope: TypedEnvelope<proto::UpdateBuffer>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::Ack> {
        this.update(&mut cx, |this, cx| {
            this.buffer_store.update(cx, |buffer_store, cx| {
                buffer_store.handle_update_buffer(envelope, false, cx)
            })
        })?
    }

    pub async fn handle_save_buffer(
        this: Model<Self>,
        envelope: TypedEnvelope<proto::SaveBuffer>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::BufferSaved> {
        let buffer_store = this.update(&mut cx, |this, _| this.buffer_store.clone())?;
        BufferStore::handle_save_buffer(buffer_store, PROJECT_ID, envelope, cx).await
    }

    pub async fn handle_open_buffer_by_path(
        this: Model<Self>,
        message: TypedEnvelope<proto::OpenBufferByPath>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::OpenBufferResponse> {
        let worktree_id = WorktreeId::from_proto(message.payload.worktree_id);
        let (buffer_store, buffer, session) = this.update(&mut cx, |this, cx| {
            let buffer_store = this.buffer_store.clone();
            let buffer = this.buffer_store.update(cx, |buffer_store, cx| {
                buffer_store.open_buffer(
                    ProjectPath {
                        worktree_id,
                        path: PathBuf::from(message.payload.path).into(),
                    },
                    cx,
                )
            });
            anyhow::Ok((buffer_store, buffer, this.session.clone()))
        })??;

        let buffer = buffer.await?;
        let buffer_id = buffer.read_with(&cx, |b, _| b.remote_id())?;

        cx.spawn(|mut cx| async move {
            BufferStore::create_buffer_for_peer(
                buffer_store,
                PEER_ID,
                buffer_id,
                PROJECT_ID,
                session,
                &mut cx,
            )
            .await
        })
        .detach();

        Ok(proto::OpenBufferResponse {
            buffer_id: buffer_id.to_proto(),
        })
    }

    pub fn on_buffer_store_event(
        &mut self,
        _: Model<BufferStore>,
        event: &BufferStoreEvent,
        cx: &mut ModelContext<Self>,
    ) {
        match event {
            BufferStoreEvent::LocalBufferUpdated { buffer } => {
                let buffer = buffer.read(cx);
                let buffer_id = buffer.remote_id();
                let Some(new_file) = buffer.file() else {
                    return;
                };
                self.session
                    .send(proto::UpdateBufferFile {
                        project_id: 0,
                        buffer_id: buffer_id.into(),
                        file: Some(new_file.to_proto(cx)),
                    })
                    .log_err();
            }
            BufferStoreEvent::DiffBaseUpdated { buffer } => {
                let buffer = buffer.read(cx);
                let buffer_id = buffer.remote_id();
                let diff_base = buffer.diff_base();
                self.session
                    .send(proto::UpdateDiffBase {
                        project_id: 0,
                        buffer_id: buffer_id.to_proto(),
                        diff_base: diff_base.map(|b| b.to_string()),
                    })
                    .log_err();
            }
            _ => {}
        }
    }
}