Merge branch 'main' into zed2-project-test

Co-authored-by: Marshall <marshall@zed.dev>

crates/Cargo.toml

@@ -0,0 +1,38 @@
+[package]
+name = "ai"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[lib]
+path = "src/ai.rs"
+doctest = false
+
+[features]
+test-support = []
+
+[dependencies]
+gpui = { path = "../gpui" }
+util = { path = "../util" }
+language = { path = "../language" }
+async-trait.workspace = true
+anyhow.workspace = true
+futures.workspace = true
+lazy_static.workspace = true
+ordered-float.workspace = true
+parking_lot.workspace = true
+isahc.workspace = true
+regex.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+postage.workspace = true
+rand.workspace = true
+log.workspace = true
+parse_duration = "2.1.1"
+tiktoken-rs = "0.5.0"
+matrixmultiply = "0.3.7"
+rusqlite = { version = "0.29.0", features = ["blob", "array", "modern_sqlite"] }
+bincode = "1.3.3"
+
+[dev-dependencies]
+gpui = { path = "../gpui", features = ["test-support"] }

crates/ai/Cargo.toml

@@ -8,9 +8,13 @@ publish = false
 path = "src/ai.rs"
 doctest = false
 
+[features]
+test-support = []
+
 [dependencies]
 gpui = { path = "../gpui" }
 util = { path = "../util" }
+language = { path = "../language" }
 async-trait.workspace = true
 anyhow.workspace = true
 futures.workspace = true

crates/ai/src/ai.rs

@@ -1,2 +1,8 @@
+pub mod auth;
 pub mod completion;
 pub mod embedding;
+pub mod models;
+pub mod prompts;
+pub mod providers;
+#[cfg(any(test, feature = "test-support"))]
+pub mod test;

crates/ai/src/auth.rs

@@ -0,0 +1,15 @@
+use gpui::AppContext;
+
+#[derive(Clone, Debug)]
+pub enum ProviderCredential {
+    Credentials { api_key: String },
+    NoCredentials,
+    NotNeeded,
+}
+
+pub trait CredentialProvider: Send + Sync {
+    fn has_credentials(&self) -> bool;
+    fn retrieve_credentials(&self, cx: &AppContext) -> ProviderCredential;
+    fn save_credentials(&self, cx: &AppContext, credential: ProviderCredential);
+    fn delete_credentials(&self, cx: &AppContext);
+}

crates/ai/src/completion.rs

@@ -1,212 +1,23 @@
-use anyhow::{anyhow, Result};
-use futures::{
-    future::BoxFuture, io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, FutureExt,
-    Stream, StreamExt,
-};
-use gpui::executor::Background;
-use isahc::{http::StatusCode, Request, RequestExt};
-use serde::{Deserialize, Serialize};
-use std::{
-    fmt::{self, Display},
-    io,
-    sync::Arc,
-};
+use anyhow::Result;
+use futures::{future::BoxFuture, stream::BoxStream};
 
-pub const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";
+use crate::{auth::CredentialProvider, models::LanguageModel};
 
-#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
-#[serde(rename_all = "lowercase")]
-pub enum Role {
-    User,
-    Assistant,
-    System,
+pub trait CompletionRequest: Send + Sync {
+    fn data(&self) -> serde_json::Result<String>;
 }
 
-impl Role {
-    pub fn cycle(&mut self) {
-        *self = match self {
-            Role::User => Role::Assistant,
-            Role::Assistant => Role::System,
-            Role::System => Role::User,
-        }
-    }
-}
-
-impl Display for Role {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Role::User => write!(f, "User"),
-            Role::Assistant => write!(f, "Assistant"),
-            Role::System => write!(f, "System"),
-        }
-    }
-}
-
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
-pub struct RequestMessage {
-    pub role: Role,
-    pub content: String,
-}
-
-#[derive(Debug, Default, Serialize)]
-pub struct OpenAIRequest {
-    pub model: String,
-    pub messages: Vec<RequestMessage>,
-    pub stream: bool,
-}
-
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
-pub struct ResponseMessage {
-    pub role: Option<Role>,
-    pub content: Option<String>,
-}
-
-#[derive(Deserialize, Debug)]
-pub struct OpenAIUsage {
-    pub prompt_tokens: u32,
-    pub completion_tokens: u32,
-    pub total_tokens: u32,
-}
-
-#[derive(Deserialize, Debug)]
-pub struct ChatChoiceDelta {
-    pub index: u32,
-    pub delta: ResponseMessage,
-    pub finish_reason: Option<String>,
-}
-
-#[derive(Deserialize, Debug)]
-pub struct OpenAIResponseStreamEvent {
-    pub id: Option<String>,
-    pub object: String,
-    pub created: u32,
-    pub model: String,
-    pub choices: Vec<ChatChoiceDelta>,
-    pub usage: Option<OpenAIUsage>,
-}
-
-pub async fn stream_completion(
-    api_key: String,
-    executor: Arc<Background>,
-    mut request: OpenAIRequest,
-) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
-    request.stream = true;
-
-    let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
-
-    let json_data = serde_json::to_string(&request)?;
-    let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
-        .header("Content-Type", "application/json")
-        .header("Authorization", format!("Bearer {}", api_key))
-        .body(json_data)?
-        .send_async()
-        .await?;
-
-    let status = response.status();
-    if status == StatusCode::OK {
-        executor
-            .spawn(async move {
-                let mut lines = BufReader::new(response.body_mut()).lines();
-
-                fn parse_line(
-                    line: Result<String, io::Error>,
-                ) -> Result<Option<OpenAIResponseStreamEvent>> {
-                    if let Some(data) = line?.strip_prefix("data: ") {
-                        let event = serde_json::from_str(&data)?;
-                        Ok(Some(event))
-                    } else {
-                        Ok(None)
-                    }
-                }
-
-                while let Some(line) = lines.next().await {
-                    if let Some(event) = parse_line(line).transpose() {
-                        let done = event.as_ref().map_or(false, |event| {
-                            event
-                                .choices
-                                .last()
-                                .map_or(false, |choice| choice.finish_reason.is_some())
-                        });
-                        if tx.unbounded_send(event).is_err() {
-                            break;
-                        }
-
-                        if done {
-                            break;
-                        }
-                    }
-                }
-
-                anyhow::Ok(())
-            })
-            .detach();
-
-        Ok(rx)
-    } else {
-        let mut body = String::new();
-        response.body_mut().read_to_string(&mut body).await?;
-
-        #[derive(Deserialize)]
-        struct OpenAIResponse {
-            error: OpenAIError,
-        }
-
-        #[derive(Deserialize)]
-        struct OpenAIError {
-            message: String,
-        }
-
-        match serde_json::from_str::<OpenAIResponse>(&body) {
-            Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
-                "Failed to connect to OpenAI API: {}",
-                response.error.message,
-            )),
-
-            _ => Err(anyhow!(
-                "Failed to connect to OpenAI API: {} {}",
-                response.status(),
-                body,
-            )),
-        }
-    }
-}
-
-pub trait CompletionProvider {
+pub trait CompletionProvider: CredentialProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel>;
     fn complete(
         &self,
-        prompt: OpenAIRequest,
+        prompt: Box<dyn CompletionRequest>,
     ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
+    fn box_clone(&self) -> Box<dyn CompletionProvider>;
 }
 
-pub struct OpenAICompletionProvider {
-    api_key: String,
-    executor: Arc<Background>,
-}
-
-impl OpenAICompletionProvider {
-    pub fn new(api_key: String, executor: Arc<Background>) -> Self {
-        Self { api_key, executor }
-    }
-}
-
-impl CompletionProvider for OpenAICompletionProvider {
-    fn complete(
-        &self,
-        prompt: OpenAIRequest,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
-        let request = stream_completion(self.api_key.clone(), self.executor.clone(), prompt);
-        async move {
-            let response = request.await?;
-            let stream = response
-                .filter_map(|response| async move {
-                    match response {
-                        Ok(mut response) => Some(Ok(response.choices.pop()?.delta.content?)),
-                        Err(error) => Some(Err(error)),
-                    }
-                })
-                .boxed();
-            Ok(stream)
-        }
-        .boxed()
+impl Clone for Box<dyn CompletionProvider> {
+    fn clone(&self) -> Box<dyn CompletionProvider> {
+        self.box_clone()
     }
 }

crates/ai/src/embedding.rs

@@ -1,30 +1,13 @@
-use anyhow::{anyhow, Result};
+use std::time::Instant;
+
+use anyhow::Result;
 use async_trait::async_trait;
-use futures::AsyncReadExt;
-use gpui::executor::Background;
-use gpui::serde_json;
-use isahc::http::StatusCode;
-use isahc::prelude::Configurable;
-use isahc::{AsyncBody, Response};
-use lazy_static::lazy_static;
 use ordered_float::OrderedFloat;
-use parking_lot::Mutex;
-use parse_duration::parse;
-use postage::watch;
 use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef};
 use rusqlite::ToSql;
-use serde::{Deserialize, Serialize};
-use std::env;
-use std::ops::Add;
-use std::sync::Arc;
-use std::time::{Duration, Instant};
-use tiktoken_rs::{cl100k_base, CoreBPE};
-use util::http::{HttpClient, Request};
 
-lazy_static! {
-    static ref OPENAI_API_KEY: Option<String> = env::var("OPENAI_API_KEY").ok();
-    static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
-}
+use crate::auth::CredentialProvider;
+use crate::models::LanguageModel;
 
 #[derive(Debug, PartialEq, Clone)]
 pub struct Embedding(pub Vec<f32>);
@@ -85,295 +68,14 @@ impl Embedding {
     }
 }
 
-// impl FromSql for Embedding {
-//     fn column_result(value: ValueRef) -> FromSqlResult<Self> {
-//         let bytes = value.as_blob()?;
-//         let embedding: Result<Vec<f32>, Box<bincode::ErrorKind>> = bincode::deserialize(bytes);
-//         if embedding.is_err() {
-//             return Err(rusqlite::types::FromSqlError::Other(embedding.unwrap_err()));
-//         }
-//         Ok(Embedding(embedding.unwrap()))
-//     }
-// }
-
-// impl ToSql for Embedding {
-//     fn to_sql(&self) -> rusqlite::Result<ToSqlOutput> {
-//         let bytes = bincode::serialize(&self.0)
-//             .map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
-//         Ok(ToSqlOutput::Owned(rusqlite::types::Value::Blob(bytes)))
-//     }
-// }
-
-#[derive(Clone)]
-pub struct OpenAIEmbeddings {
-    pub client: Arc<dyn HttpClient>,
-    pub executor: Arc<Background>,
-    rate_limit_count_rx: watch::Receiver<Option<Instant>>,
-    rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
-}
-
-#[derive(Serialize)]
-struct OpenAIEmbeddingRequest<'a> {
-    model: &'static str,
-    input: Vec<&'a str>,
-}
-
-#[derive(Deserialize)]
-struct OpenAIEmbeddingResponse {
-    data: Vec<OpenAIEmbedding>,
-    usage: OpenAIEmbeddingUsage,
-}
-
-#[derive(Debug, Deserialize)]
-struct OpenAIEmbedding {
-    embedding: Vec<f32>,
-    index: usize,
-    object: String,
-}
-
-#[derive(Deserialize)]
-struct OpenAIEmbeddingUsage {
-    prompt_tokens: usize,
-    total_tokens: usize,
-}
-
 #[async_trait]
-pub trait EmbeddingProvider: Sync + Send {
-    fn is_authenticated(&self) -> bool;
+pub trait EmbeddingProvider: CredentialProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel>;
     async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
     fn max_tokens_per_batch(&self) -> usize;
-    fn truncate(&self, span: &str) -> (String, usize);
     fn rate_limit_expiration(&self) -> Option<Instant>;
 }
 
-pub struct DummyEmbeddings {}
-
-#[async_trait]
-impl EmbeddingProvider for DummyEmbeddings {
-    fn is_authenticated(&self) -> bool {
-        true
-    }
-    fn rate_limit_expiration(&self) -> Option<Instant> {
-        None
-    }
-    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
-        // 1024 is the OpenAI Embeddings size for ada models.
-        // the model we will likely be starting with.
-        let dummy_vec = Embedding::from(vec![0.32 as f32; 1536]);
-        return Ok(vec![dummy_vec; spans.len()]);
-    }
-
-    fn max_tokens_per_batch(&self) -> usize {
-        OPENAI_INPUT_LIMIT
-    }
-
-    fn truncate(&self, span: &str) -> (String, usize) {
-        let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
-        let token_count = tokens.len();
-        let output = if token_count > OPENAI_INPUT_LIMIT {
-            tokens.truncate(OPENAI_INPUT_LIMIT);
-            let new_input = OPENAI_BPE_TOKENIZER.decode(tokens.clone());
-            new_input.ok().unwrap_or_else(|| span.to_string())
-        } else {
-            span.to_string()
-        };
-
-        (output, tokens.len())
-    }
-}
-
-const OPENAI_INPUT_LIMIT: usize = 8190;
-
-impl OpenAIEmbeddings {
-    pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Background>) -> Self {
-        let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
-        let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
-
-        OpenAIEmbeddings {
-            client,
-            executor,
-            rate_limit_count_rx,
-            rate_limit_count_tx,
-        }
-    }
-
-    fn resolve_rate_limit(&self) {
-        let reset_time = *self.rate_limit_count_tx.lock().borrow();
-
-        if let Some(reset_time) = reset_time {
-            if Instant::now() >= reset_time {
-                *self.rate_limit_count_tx.lock().borrow_mut() = None
-            }
-        }
-
-        log::trace!(
-            "resolving reset time: {:?}",
-            *self.rate_limit_count_tx.lock().borrow()
-        );
-    }
-
-    fn update_reset_time(&self, reset_time: Instant) {
-        let original_time = *self.rate_limit_count_tx.lock().borrow();
-
-        let updated_time = if let Some(original_time) = original_time {
-            if reset_time < original_time {
-                Some(reset_time)
-            } else {
-                Some(original_time)
-            }
-        } else {
-            Some(reset_time)
-        };
-
-        log::trace!("updating rate limit time: {:?}", updated_time);
-
-        *self.rate_limit_count_tx.lock().borrow_mut() = updated_time;
-    }
-    async fn send_request(
-        &self,
-        api_key: &str,
-        spans: Vec<&str>,
-        request_timeout: u64,
-    ) -> Result<Response<AsyncBody>> {
-        let request = Request::post("https://api.openai.com/v1/embeddings")
-            .redirect_policy(isahc::config::RedirectPolicy::Follow)
-            .timeout(Duration::from_secs(request_timeout))
-            .header("Content-Type", "application/json")
-            .header("Authorization", format!("Bearer {}", api_key))
-            .body(
-                serde_json::to_string(&OpenAIEmbeddingRequest {
-                    input: spans.clone(),
-                    model: "text-embedding-ada-002",
-                })
-                .unwrap()
-                .into(),
-            )?;
-
-        Ok(self.client.send(request).await?)
-    }
-}
-
-#[async_trait]
-impl EmbeddingProvider for OpenAIEmbeddings {
-    fn is_authenticated(&self) -> bool {
-        OPENAI_API_KEY.as_ref().is_some()
-    }
-    fn max_tokens_per_batch(&self) -> usize {
-        50000
-    }
-
-    fn rate_limit_expiration(&self) -> Option<Instant> {
-        *self.rate_limit_count_rx.borrow()
-    }
-    fn truncate(&self, span: &str) -> (String, usize) {
-        let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
-        let output = if tokens.len() > OPENAI_INPUT_LIMIT {
-            tokens.truncate(OPENAI_INPUT_LIMIT);
-            OPENAI_BPE_TOKENIZER
-                .decode(tokens.clone())
-                .ok()
-                .unwrap_or_else(|| span.to_string())
-        } else {
-            span.to_string()
-        };
-
-        (output, tokens.len())
-    }
-
-    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
-        const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
-        const MAX_RETRIES: usize = 4;
-
-        let api_key = OPENAI_API_KEY
-            .as_ref()
-            .ok_or_else(|| anyhow!("no api key"))?;
-
-        let mut request_number = 0;
-        let mut rate_limiting = false;
-        let mut request_timeout: u64 = 15;
-        let mut response: Response<AsyncBody>;
-        while request_number < MAX_RETRIES {
-            response = self
-                .send_request(
-                    api_key,
-                    spans.iter().map(|x| &**x).collect(),
-                    request_timeout,
-                )
-                .await?;
-            request_number += 1;
-
-            match response.status() {
-                StatusCode::REQUEST_TIMEOUT => {
-                    request_timeout += 5;
-                }
-                StatusCode::OK => {
-                    let mut body = String::new();
-                    response.body_mut().read_to_string(&mut body).await?;
-                    let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
-
-                    log::trace!(
-                        "openai embedding completed. tokens: {:?}",
-                        response.usage.total_tokens
-                    );
-
-                    // If we complete a request successfully that was previously rate_limited
-                    // resolve the rate limit
-                    if rate_limiting {
-                        self.resolve_rate_limit()
-                    }
-
-                    return Ok(response
-                        .data
-                        .into_iter()
-                        .map(|embedding| Embedding::from(embedding.embedding))
-                        .collect());
-                }
-                StatusCode::TOO_MANY_REQUESTS => {
-                    rate_limiting = true;
-                    let mut body = String::new();
-                    response.body_mut().read_to_string(&mut body).await?;
-
-                    let delay_duration = {
-                        let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
-                        if let Some(time_to_reset) =
-                            response.headers().get("x-ratelimit-reset-tokens")
-                        {
-                            if let Ok(time_str) = time_to_reset.to_str() {
-                                parse(time_str).unwrap_or(delay)
-                            } else {
-                                delay
-                            }
-                        } else {
-                            delay
-                        }
-                    };
-
-                    // If we've previously rate limited, increment the duration but not the count
-                    let reset_time = Instant::now().add(delay_duration);
-                    self.update_reset_time(reset_time);
-
-                    log::trace!(
-                        "openai rate limiting: waiting {:?} until lifted",
-                        &delay_duration
-                    );
-
-                    self.executor.timer(delay_duration).await;
-                }
-                _ => {
-                    let mut body = String::new();
-                    response.body_mut().read_to_string(&mut body).await?;
-                    return Err(anyhow!(
-                        "open ai bad request: {:?} {:?}",
-                        &response.status(),
-                        body
-                    ));
-                }
-            }
-        }
-        Err(anyhow!("openai max retries"))
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

crates/ai/src/models.rs

@@ -0,0 +1,16 @@
+pub enum TruncationDirection {
+    Start,
+    End,
+}
+
+pub trait LanguageModel {
+    fn name(&self) -> String;
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize>;
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String>;
+    fn capacity(&self) -> anyhow::Result<usize>;
+}

crates/ai/src/prompts/base.rs

@@ -0,0 +1,330 @@
+use std::cmp::Reverse;
+use std::ops::Range;
+use std::sync::Arc;
+
+use language::BufferSnapshot;
+use util::ResultExt;
+
+use crate::models::LanguageModel;
+use crate::prompts::repository_context::PromptCodeSnippet;
+
+pub(crate) enum PromptFileType {
+    Text,
+    Code,
+}
+
+// TODO: Set this up to manage for defaults well
+pub struct PromptArguments {
+    pub model: Arc<dyn LanguageModel>,
+    pub user_prompt: Option<String>,
+    pub language_name: Option<String>,
+    pub project_name: Option<String>,
+    pub snippets: Vec<PromptCodeSnippet>,
+    pub reserved_tokens: usize,
+    pub buffer: Option<BufferSnapshot>,
+    pub selected_range: Option<Range<usize>>,
+}
+
+impl PromptArguments {
+    pub(crate) fn get_file_type(&self) -> PromptFileType {
+        if self
+            .language_name
+            .as_ref()
+            .and_then(|name| Some(!["Markdown", "Plain Text"].contains(&name.as_str())))
+            .unwrap_or(true)
+        {
+            PromptFileType::Code
+        } else {
+            PromptFileType::Text
+        }
+    }
+}
+
+pub trait PromptTemplate {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)>;
+}
+
+#[repr(i8)]
+#[derive(PartialEq, Eq, Ord)]
+pub enum PromptPriority {
+    Mandatory,                // Ignores truncation
+    Ordered { order: usize }, // Truncates based on priority
+}
+
+impl PartialOrd for PromptPriority {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        match (self, other) {
+            (Self::Mandatory, Self::Mandatory) => Some(std::cmp::Ordering::Equal),
+            (Self::Mandatory, Self::Ordered { .. }) => Some(std::cmp::Ordering::Greater),
+            (Self::Ordered { .. }, Self::Mandatory) => Some(std::cmp::Ordering::Less),
+            (Self::Ordered { order: a }, Self::Ordered { order: b }) => b.partial_cmp(a),
+        }
+    }
+}
+
+pub struct PromptChain {
+    args: PromptArguments,
+    templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
+}
+
+impl PromptChain {
+    pub fn new(
+        args: PromptArguments,
+        templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
+    ) -> Self {
+        PromptChain { args, templates }
+    }
+
+    pub fn generate(&self, truncate: bool) -> anyhow::Result<(String, usize)> {
+        // Argsort based on Prompt Priority
+        let seperator = "\n";
+        let seperator_tokens = self.args.model.count_tokens(seperator)?;
+        let mut sorted_indices = (0..self.templates.len()).collect::<Vec<_>>();
+        sorted_indices.sort_by_key(|&i| Reverse(&self.templates[i].0));
+
+        // If Truncate
+        let mut tokens_outstanding = if truncate {
+            Some(self.args.model.capacity()? - self.args.reserved_tokens)
+        } else {
+            None
+        };
+
+        let mut prompts = vec!["".to_string(); sorted_indices.len()];
+        for idx in sorted_indices {
+            let (_, template) = &self.templates[idx];
+
+            if let Some((template_prompt, prompt_token_count)) =
+                template.generate(&self.args, tokens_outstanding).log_err()
+            {
+                if template_prompt != "" {
+                    prompts[idx] = template_prompt;
+
+                    if let Some(remaining_tokens) = tokens_outstanding {
+                        let new_tokens = prompt_token_count + seperator_tokens;
+                        tokens_outstanding = if remaining_tokens > new_tokens {
+                            Some(remaining_tokens - new_tokens)
+                        } else {
+                            Some(0)
+                        };
+                    }
+                }
+            }
+        }
+
+        prompts.retain(|x| x != "");
+
+        let full_prompt = prompts.join(seperator);
+        let total_token_count = self.args.model.count_tokens(&full_prompt)?;
+        anyhow::Ok((prompts.join(seperator), total_token_count))
+    }
+}
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use crate::models::TruncationDirection;
+    use crate::test::FakeLanguageModel;
+
+    use super::*;
+
+    #[test]
+    pub fn test_prompt_chain() {
+        struct TestPromptTemplate {}
+        impl PromptTemplate for TestPromptTemplate {
+            fn generate(
+                &self,
+                args: &PromptArguments,
+                max_token_length: Option<usize>,
+            ) -> anyhow::Result<(String, usize)> {
+                let mut content = "This is a test prompt template".to_string();
+
+                let mut token_count = args.model.count_tokens(&content)?;
+                if let Some(max_token_length) = max_token_length {
+                    if token_count > max_token_length {
+                        content = args.model.truncate(
+                            &content,
+                            max_token_length,
+                            TruncationDirection::End,
+                        )?;
+                        token_count = max_token_length;
+                    }
+                }
+
+                anyhow::Ok((content, token_count))
+            }
+        }
+
+        struct TestLowPriorityTemplate {}
+        impl PromptTemplate for TestLowPriorityTemplate {
+            fn generate(
+                &self,
+                args: &PromptArguments,
+                max_token_length: Option<usize>,
+            ) -> anyhow::Result<(String, usize)> {
+                let mut content = "This is a low priority test prompt template".to_string();
+
+                let mut token_count = args.model.count_tokens(&content)?;
+                if let Some(max_token_length) = max_token_length {
+                    if token_count > max_token_length {
+                        content = args.model.truncate(
+                            &content,
+                            max_token_length,
+                            TruncationDirection::End,
+                        )?;
+                        token_count = max_token_length;
+                    }
+                }
+
+                anyhow::Ok((content, token_count))
+            }
+        }
+
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity: 100 });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(false).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a test prompt template\nThis is a low priority test prompt template"
+                .to_string()
+        );
+
+        assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
+
+        // Testing with Truncation Off
+        // Should ignore capacity and return all prompts
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity: 20 });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(false).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a test prompt template\nThis is a low priority test prompt template"
+                .to_string()
+        );
+
+        assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
+
+        // Testing with Truncation Off
+        // Should ignore capacity and return all prompts
+        let capacity = 20;
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 2 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(true).unwrap();
+
+        assert_eq!(prompt, "This is a test promp".to_string());
+        assert_eq!(token_count, capacity);
+
+        // Change Ordering of Prompts Based on Priority
+        let capacity = 120;
+        let reserved_tokens = 10;
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Mandatory,
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(true).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a low priority test prompt template\nThis is a test prompt template\nThis is a low priority test prompt "
+                .to_string()
+        );
+        assert_eq!(token_count, capacity - reserved_tokens);
+    }
+}

crates/ai/src/prompts/file_context.rs

@@ -0,0 +1,164 @@
+use anyhow::anyhow;
+use language::BufferSnapshot;
+use language::ToOffset;
+
+use crate::models::LanguageModel;
+use crate::models::TruncationDirection;
+use crate::prompts::base::PromptArguments;
+use crate::prompts::base::PromptTemplate;
+use std::fmt::Write;
+use std::ops::Range;
+use std::sync::Arc;
+
+fn retrieve_context(
+    buffer: &BufferSnapshot,
+    selected_range: &Option<Range<usize>>,
+    model: Arc<dyn LanguageModel>,
+    max_token_count: Option<usize>,
+) -> anyhow::Result<(String, usize, bool)> {
+    let mut prompt = String::new();
+    let mut truncated = false;
+    if let Some(selected_range) = selected_range {
+        let start = selected_range.start.to_offset(buffer);
+        let end = selected_range.end.to_offset(buffer);
+
+        let start_window = buffer.text_for_range(0..start).collect::<String>();
+
+        let mut selected_window = String::new();
+        if start == end {
+            write!(selected_window, "<|START|>").unwrap();
+        } else {
+            write!(selected_window, "<|START|").unwrap();
+        }
+
+        write!(
+            selected_window,
+            "{}",
+            buffer.text_for_range(start..end).collect::<String>()
+        )
+        .unwrap();
+
+        if start != end {
+            write!(selected_window, "|END|>").unwrap();
+        }
+
+        let end_window = buffer.text_for_range(end..buffer.len()).collect::<String>();
+
+        if let Some(max_token_count) = max_token_count {
+            let selected_tokens = model.count_tokens(&selected_window)?;
+            if selected_tokens > max_token_count {
+                return Err(anyhow!(
+                    "selected range is greater than model context window, truncation not possible"
+                ));
+            };
+
+            let mut remaining_tokens = max_token_count - selected_tokens;
+            let start_window_tokens = model.count_tokens(&start_window)?;
+            let end_window_tokens = model.count_tokens(&end_window)?;
+            let outside_tokens = start_window_tokens + end_window_tokens;
+            if outside_tokens > remaining_tokens {
+                let (start_goal_tokens, end_goal_tokens) =
+                    if start_window_tokens < end_window_tokens {
+                        let start_goal_tokens = (remaining_tokens / 2).min(start_window_tokens);
+                        remaining_tokens -= start_goal_tokens;
+                        let end_goal_tokens = remaining_tokens.min(end_window_tokens);
+                        (start_goal_tokens, end_goal_tokens)
+                    } else {
+                        let end_goal_tokens = (remaining_tokens / 2).min(end_window_tokens);
+                        remaining_tokens -= end_goal_tokens;
+                        let start_goal_tokens = remaining_tokens.min(start_window_tokens);
+                        (start_goal_tokens, end_goal_tokens)
+                    };
+
+                let truncated_start_window =
+                    model.truncate(&start_window, start_goal_tokens, TruncationDirection::Start)?;
+                let truncated_end_window =
+                    model.truncate(&end_window, end_goal_tokens, TruncationDirection::End)?;
+                writeln!(
+                    prompt,
+                    "{truncated_start_window}{selected_window}{truncated_end_window}"
+                )
+                .unwrap();
+                truncated = true;
+            } else {
+                writeln!(prompt, "{start_window}{selected_window}{end_window}").unwrap();
+            }
+        } else {
+            // If we dont have a selected range, include entire file.
+            writeln!(prompt, "{}", &buffer.text()).unwrap();
+
+            // Dumb truncation strategy
+            if let Some(max_token_count) = max_token_count {
+                if model.count_tokens(&prompt)? > max_token_count {
+                    truncated = true;
+                    prompt = model.truncate(&prompt, max_token_count, TruncationDirection::End)?;
+                }
+            }
+        }
+    }
+
+    let token_count = model.count_tokens(&prompt)?;
+    anyhow::Ok((prompt, token_count, truncated))
+}
+
+pub struct FileContext {}
+
+impl PromptTemplate for FileContext {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        if let Some(buffer) = &args.buffer {
+            let mut prompt = String::new();
+            // Add Initial Preamble
+            // TODO: Do we want to add the path in here?
+            writeln!(
+                prompt,
+                "The file you are currently working on has the following content:"
+            )
+            .unwrap();
+
+            let language_name = args
+                .language_name
+                .clone()
+                .unwrap_or("".to_string())
+                .to_lowercase();
+
+            let (context, _, truncated) = retrieve_context(
+                buffer,
+                &args.selected_range,
+                args.model.clone(),
+                max_token_length,
+            )?;
+            writeln!(prompt, "```{language_name}\n{context}\n```").unwrap();
+
+            if truncated {
+                writeln!(prompt, "Note the content has been truncated and only represents a portion of the file.").unwrap();
+            }
+
+            if let Some(selected_range) = &args.selected_range {
+                let start = selected_range.start.to_offset(buffer);
+                let end = selected_range.end.to_offset(buffer);
+
+                if start == end {
+                    writeln!(prompt, "In particular, the user's cursor is currently on the '<|START|>' span in the above content, with no text selected.").unwrap();
+                } else {
+                    writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
+                }
+            }
+
+            // Really dumb truncation strategy
+            if let Some(max_tokens) = max_token_length {
+                prompt = args
+                    .model
+                    .truncate(&prompt, max_tokens, TruncationDirection::End)?;
+            }
+
+            let token_count = args.model.count_tokens(&prompt)?;
+            anyhow::Ok((prompt, token_count))
+        } else {
+            Err(anyhow!("no buffer provided to retrieve file context from"))
+        }
+    }
+}

crates/ai/src/prompts/generate.rs

@@ -0,0 +1,99 @@
+use crate::prompts::base::{PromptArguments, PromptFileType, PromptTemplate};
+use anyhow::anyhow;
+use std::fmt::Write;
+
+pub fn capitalize(s: &str) -> String {
+    let mut c = s.chars();
+    match c.next() {
+        None => String::new(),
+        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
+    }
+}
+
+pub struct GenerateInlineContent {}
+
+impl PromptTemplate for GenerateInlineContent {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let Some(user_prompt) = &args.user_prompt else {
+            return Err(anyhow!("user prompt not provided"));
+        };
+
+        let file_type = args.get_file_type();
+        let content_type = match &file_type {
+            PromptFileType::Code => "code",
+            PromptFileType::Text => "text",
+        };
+
+        let mut prompt = String::new();
+
+        if let Some(selected_range) = &args.selected_range {
+            if selected_range.start == selected_range.end {
+                writeln!(
+                    prompt,
+                    "Assume the cursor is located where the `<|START|>` span is."
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "{} can't be replaced, so assume your answer will be inserted at the cursor.",
+                    capitalize(content_type)
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "Generate {content_type} based on the users prompt: {user_prompt}",
+                )
+                .unwrap();
+            } else {
+                writeln!(prompt, "Modify the user's selected {content_type} based upon the users prompt: '{user_prompt}'").unwrap();
+                writeln!(prompt, "You must reply with only the adjusted {content_type} (within the '<|START|' and '|END|>' spans) not the entire file.").unwrap();
+                writeln!(prompt, "Double check that you only return code and not the '<|START|' and '|END|'> spans").unwrap();
+            }
+        } else {
+            writeln!(
+                prompt,
+                "Generate {content_type} based on the users prompt: {user_prompt}"
+            )
+            .unwrap();
+        }
+
+        if let Some(language_name) = &args.language_name {
+            writeln!(
+                prompt,
+                "Your answer MUST always and only be valid {}.",
+                language_name
+            )
+            .unwrap();
+        }
+        writeln!(prompt, "Never make remarks about the output.").unwrap();
+        writeln!(
+            prompt,
+            "Do not return anything else, except the generated {content_type}."
+        )
+        .unwrap();
+
+        match file_type {
+            PromptFileType::Code => {
+                // writeln!(prompt, "Always wrap your code in a Markdown block.").unwrap();
+            }
+            _ => {}
+        }
+
+        // Really dumb truncation strategy
+        if let Some(max_tokens) = max_token_length {
+            prompt = args.model.truncate(
+                &prompt,
+                max_tokens,
+                crate::models::TruncationDirection::End,
+            )?;
+        }
+
+        let token_count = args.model.count_tokens(&prompt)?;
+
+        anyhow::Ok((prompt, token_count))
+    }
+}

crates/ai/src/prompts/mod.rs

@@ -0,0 +1,5 @@
+pub mod base;
+pub mod file_context;
+pub mod generate;
+pub mod preamble;
+pub mod repository_context;

crates/ai/src/prompts/preamble.rs

@@ -0,0 +1,52 @@
+use crate::prompts::base::{PromptArguments, PromptFileType, PromptTemplate};
+use std::fmt::Write;
+
+pub struct EngineerPreamble {}
+
+impl PromptTemplate for EngineerPreamble {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let mut prompts = Vec::new();
+
+        match args.get_file_type() {
+            PromptFileType::Code => {
+                prompts.push(format!(
+                    "You are an expert {}engineer.",
+                    args.language_name.clone().unwrap_or("".to_string()) + " "
+                ));
+            }
+            PromptFileType::Text => {
+                prompts.push("You are an expert engineer.".to_string());
+            }
+        }
+
+        if let Some(project_name) = args.project_name.clone() {
+            prompts.push(format!(
+                "You are currently working inside the '{project_name}' project in code editor Zed."
+            ));
+        }
+
+        if let Some(mut remaining_tokens) = max_token_length {
+            let mut prompt = String::new();
+            let mut total_count = 0;
+            for prompt_piece in prompts {
+                let prompt_token_count =
+                    args.model.count_tokens(&prompt_piece)? + args.model.count_tokens("\n")?;
+                if remaining_tokens > prompt_token_count {
+                    writeln!(prompt, "{prompt_piece}").unwrap();
+                    remaining_tokens -= prompt_token_count;
+                    total_count += prompt_token_count;
+                }
+            }
+
+            anyhow::Ok((prompt, total_count))
+        } else {
+            let prompt = prompts.join("\n");
+            let token_count = args.model.count_tokens(&prompt)?;
+            anyhow::Ok((prompt, token_count))
+        }
+    }
+}

crates/ai/src/prompts/repository_context.rs

@@ -0,0 +1,94 @@
+use crate::prompts::base::{PromptArguments, PromptTemplate};
+use std::fmt::Write;
+use std::{ops::Range, path::PathBuf};
+
+use gpui::{AsyncAppContext, ModelHandle};
+use language::{Anchor, Buffer};
+
+#[derive(Clone)]
+pub struct PromptCodeSnippet {
+    path: Option<PathBuf>,
+    language_name: Option<String>,
+    content: String,
+}
+
+impl PromptCodeSnippet {
+    pub fn new(buffer: ModelHandle<Buffer>, range: Range<Anchor>, cx: &AsyncAppContext) -> Self {
+        let (content, language_name, file_path) = buffer.read_with(cx, |buffer, _| {
+            let snapshot = buffer.snapshot();
+            let content = snapshot.text_for_range(range.clone()).collect::<String>();
+
+            let language_name = buffer
+                .language()
+                .and_then(|language| Some(language.name().to_string().to_lowercase()));
+
+            let file_path = buffer
+                .file()
+                .and_then(|file| Some(file.path().to_path_buf()));
+
+            (content, language_name, file_path)
+        });
+
+        PromptCodeSnippet {
+            path: file_path,
+            language_name,
+            content,
+        }
+    }
+}
+
+impl ToString for PromptCodeSnippet {
+    fn to_string(&self) -> String {
+        let path = self
+            .path
+            .as_ref()
+            .and_then(|path| Some(path.to_string_lossy().to_string()))
+            .unwrap_or("".to_string());
+        let language_name = self.language_name.clone().unwrap_or("".to_string());
+        let content = self.content.clone();
+
+        format!("The below code snippet may be relevant from file: {path}\n```{language_name}\n{content}\n```")
+    }
+}
+
+pub struct RepositoryContext {}
+
+impl PromptTemplate for RepositoryContext {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        const MAXIMUM_SNIPPET_TOKEN_COUNT: usize = 500;
+        let template = "You are working inside a large repository, here are a few code snippets that may be useful.";
+        let mut prompt = String::new();
+
+        let mut remaining_tokens = max_token_length.clone();
+        let seperator_token_length = args.model.count_tokens("\n")?;
+        for snippet in &args.snippets {
+            let mut snippet_prompt = template.to_string();
+            let content = snippet.to_string();
+            writeln!(snippet_prompt, "{content}").unwrap();
+
+            let token_count = args.model.count_tokens(&snippet_prompt)?;
+            if token_count <= MAXIMUM_SNIPPET_TOKEN_COUNT {
+                if let Some(tokens_left) = remaining_tokens {
+                    if tokens_left >= token_count {
+                        writeln!(prompt, "{snippet_prompt}").unwrap();
+                        remaining_tokens = if tokens_left >= (token_count + seperator_token_length)
+                        {
+                            Some(tokens_left - token_count - seperator_token_length)
+                        } else {
+                            Some(0)
+                        };
+                    }
+                } else {
+                    writeln!(prompt, "{snippet_prompt}").unwrap();
+                }
+            }
+        }
+
+        let total_token_count = args.model.count_tokens(&prompt)?;
+        anyhow::Ok((prompt, total_token_count))
+    }
+}

crates/ai/src/providers/mod.rs

@@ -0,0 +1 @@
+pub mod open_ai;

crates/ai/src/providers/open_ai/completion.rs

@@ -0,0 +1,298 @@
+use anyhow::{anyhow, Result};
+use futures::{
+    future::BoxFuture, io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, FutureExt,
+    Stream, StreamExt,
+};
+use gpui::{executor::Background, AppContext};
+use isahc::{http::StatusCode, Request, RequestExt};
+use parking_lot::RwLock;
+use serde::{Deserialize, Serialize};
+use std::{
+    env,
+    fmt::{self, Display},
+    io,
+    sync::Arc,
+};
+use util::ResultExt;
+
+use crate::{
+    auth::{CredentialProvider, ProviderCredential},
+    completion::{CompletionProvider, CompletionRequest},
+    models::LanguageModel,
+};
+
+use crate::providers::open_ai::{OpenAILanguageModel, OPENAI_API_URL};
+
+#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    User,
+    Assistant,
+    System,
+}
+
+impl Role {
+    pub fn cycle(&mut self) {
+        *self = match self {
+            Role::User => Role::Assistant,
+            Role::Assistant => Role::System,
+            Role::System => Role::User,
+        }
+    }
+}
+
+impl Display for Role {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Role::User => write!(f, "User"),
+            Role::Assistant => write!(f, "Assistant"),
+            Role::System => write!(f, "System"),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+pub struct RequestMessage {
+    pub role: Role,
+    pub content: String,
+}
+
+#[derive(Debug, Default, Serialize)]
+pub struct OpenAIRequest {
+    pub model: String,
+    pub messages: Vec<RequestMessage>,
+    pub stream: bool,
+    pub stop: Vec<String>,
+    pub temperature: f32,
+}
+
+impl CompletionRequest for OpenAIRequest {
+    fn data(&self) -> serde_json::Result<String> {
+        serde_json::to_string(self)
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+pub struct ResponseMessage {
+    pub role: Option<Role>,
+    pub content: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct OpenAIUsage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct ChatChoiceDelta {
+    pub index: u32,
+    pub delta: ResponseMessage,
+    pub finish_reason: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct OpenAIResponseStreamEvent {
+    pub id: Option<String>,
+    pub object: String,
+    pub created: u32,
+    pub model: String,
+    pub choices: Vec<ChatChoiceDelta>,
+    pub usage: Option<OpenAIUsage>,
+}
+
+pub async fn stream_completion(
+    credential: ProviderCredential,
+    executor: Arc<Background>,
+    request: Box<dyn CompletionRequest>,
+) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
+    let api_key = match credential {
+        ProviderCredential::Credentials { api_key } => api_key,
+        _ => {
+            return Err(anyhow!("no credentials provider for completion"));
+        }
+    };
+
+    let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
+
+    let json_data = request.data()?;
+    let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
+        .header("Content-Type", "application/json")
+        .header("Authorization", format!("Bearer {}", api_key))
+        .body(json_data)?
+        .send_async()
+        .await?;
+
+    let status = response.status();
+    if status == StatusCode::OK {
+        executor
+            .spawn(async move {
+                let mut lines = BufReader::new(response.body_mut()).lines();
+
+                fn parse_line(
+                    line: Result<String, io::Error>,
+                ) -> Result<Option<OpenAIResponseStreamEvent>> {
+                    if let Some(data) = line?.strip_prefix("data: ") {
+                        let event = serde_json::from_str(&data)?;
+                        Ok(Some(event))
+                    } else {
+                        Ok(None)
+                    }
+                }
+
+                while let Some(line) = lines.next().await {
+                    if let Some(event) = parse_line(line).transpose() {
+                        let done = event.as_ref().map_or(false, |event| {
+                            event
+                                .choices
+                                .last()
+                                .map_or(false, |choice| choice.finish_reason.is_some())
+                        });
+                        if tx.unbounded_send(event).is_err() {
+                            break;
+                        }
+
+                        if done {
+                            break;
+                        }
+                    }
+                }
+
+                anyhow::Ok(())
+            })
+            .detach();
+
+        Ok(rx)
+    } else {
+        let mut body = String::new();
+        response.body_mut().read_to_string(&mut body).await?;
+
+        #[derive(Deserialize)]
+        struct OpenAIResponse {
+            error: OpenAIError,
+        }
+
+        #[derive(Deserialize)]
+        struct OpenAIError {
+            message: String,
+        }
+
+        match serde_json::from_str::<OpenAIResponse>(&body) {
+            Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
+                "Failed to connect to OpenAI API: {}",
+                response.error.message,
+            )),
+
+            _ => Err(anyhow!(
+                "Failed to connect to OpenAI API: {} {}",
+                response.status(),
+                body,
+            )),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct OpenAICompletionProvider {
+    model: OpenAILanguageModel,
+    credential: Arc<RwLock<ProviderCredential>>,
+    executor: Arc<Background>,
+}
+
+impl OpenAICompletionProvider {
+    pub fn new(model_name: &str, executor: Arc<Background>) -> Self {
+        let model = OpenAILanguageModel::load(model_name);
+        let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
+        Self {
+            model,
+            credential,
+            executor,
+        }
+    }
+}
+
+impl CredentialProvider for OpenAICompletionProvider {
+    fn has_credentials(&self) -> bool {
+        match *self.credential.read() {
+            ProviderCredential::Credentials { .. } => true,
+            _ => false,
+        }
+    }
+    fn retrieve_credentials(&self, cx: &AppContext) -> ProviderCredential {
+        let mut credential = self.credential.write();
+        match *credential {
+            ProviderCredential::Credentials { .. } => {
+                return credential.clone();
+            }
+            _ => {
+                if let Ok(api_key) = env::var("OPENAI_API_KEY") {
+                    *credential = ProviderCredential::Credentials { api_key };
+                } else if let Some((_, api_key)) = cx
+                    .platform()
+                    .read_credentials(OPENAI_API_URL)
+                    .log_err()
+                    .flatten()
+                {
+                    if let Some(api_key) = String::from_utf8(api_key).log_err() {
+                        *credential = ProviderCredential::Credentials { api_key };
+                    }
+                } else {
+                };
+            }
+        }
+
+        credential.clone()
+    }
+
+    fn save_credentials(&self, cx: &AppContext, credential: ProviderCredential) {
+        match credential.clone() {
+            ProviderCredential::Credentials { api_key } => {
+                cx.platform()
+                    .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
+                    .log_err();
+            }
+            _ => {}
+        }
+
+        *self.credential.write() = credential;
+    }
+    fn delete_credentials(&self, cx: &AppContext) {
+        cx.platform().delete_credentials(OPENAI_API_URL).log_err();
+        *self.credential.write() = ProviderCredential::NoCredentials;
+    }
+}
+
+impl CompletionProvider for OpenAICompletionProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(self.model.clone());
+        model
+    }
+    fn complete(
+        &self,
+        prompt: Box<dyn CompletionRequest>,
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
+        // Currently the CompletionRequest for OpenAI, includes a 'model' parameter
+        // This means that the model is determined by the CompletionRequest and not the CompletionProvider,
+        // which is currently model based, due to the langauge model.
+        // At some point in the future we should rectify this.
+        let credential = self.credential.read().clone();
+        let request = stream_completion(credential, self.executor.clone(), prompt);
+        async move {
+            let response = request.await?;
+            let stream = response
+                .filter_map(|response| async move {
+                    match response {
+                        Ok(mut response) => Some(Ok(response.choices.pop()?.delta.content?)),
+                        Err(error) => Some(Err(error)),
+                    }
+                })
+                .boxed();
+            Ok(stream)
+        }
+        .boxed()
+    }
+    fn box_clone(&self) -> Box<dyn CompletionProvider> {
+        Box::new((*self).clone())
+    }
+}

crates/ai/src/providers/open_ai/embedding.rs

@@ -0,0 +1,306 @@
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use futures::AsyncReadExt;
+use gpui::executor::Background;
+use gpui::{serde_json, AppContext};
+use isahc::http::StatusCode;
+use isahc::prelude::Configurable;
+use isahc::{AsyncBody, Response};
+use lazy_static::lazy_static;
+use parking_lot::{Mutex, RwLock};
+use parse_duration::parse;
+use postage::watch;
+use serde::{Deserialize, Serialize};
+use std::env;
+use std::ops::Add;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tiktoken_rs::{cl100k_base, CoreBPE};
+use util::http::{HttpClient, Request};
+use util::ResultExt;
+
+use crate::auth::{CredentialProvider, ProviderCredential};
+use crate::embedding::{Embedding, EmbeddingProvider};
+use crate::models::LanguageModel;
+use crate::providers::open_ai::OpenAILanguageModel;
+
+use crate::providers::open_ai::OPENAI_API_URL;
+
+lazy_static! {
+    static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
+}
+
+#[derive(Clone)]
+pub struct OpenAIEmbeddingProvider {
+    model: OpenAILanguageModel,
+    credential: Arc<RwLock<ProviderCredential>>,
+    pub client: Arc<dyn HttpClient>,
+    pub executor: Arc<Background>,
+    rate_limit_count_rx: watch::Receiver<Option<Instant>>,
+    rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
+}
+
+#[derive(Serialize)]
+struct OpenAIEmbeddingRequest<'a> {
+    model: &'static str,
+    input: Vec<&'a str>,
+}
+
+#[derive(Deserialize)]
+struct OpenAIEmbeddingResponse {
+    data: Vec<OpenAIEmbedding>,
+    usage: OpenAIEmbeddingUsage,
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenAIEmbedding {
+    embedding: Vec<f32>,
+    index: usize,
+    object: String,
+}
+
+#[derive(Deserialize)]
+struct OpenAIEmbeddingUsage {
+    prompt_tokens: usize,
+    total_tokens: usize,
+}
+
+impl OpenAIEmbeddingProvider {
+    pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Background>) -> Self {
+        let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
+        let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
+
+        let model = OpenAILanguageModel::load("text-embedding-ada-002");
+        let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
+
+        OpenAIEmbeddingProvider {
+            model,
+            credential,
+            client,
+            executor,
+            rate_limit_count_rx,
+            rate_limit_count_tx,
+        }
+    }
+
+    fn get_api_key(&self) -> Result<String> {
+        match self.credential.read().clone() {
+            ProviderCredential::Credentials { api_key } => Ok(api_key),
+            _ => Err(anyhow!("api credentials not provided")),
+        }
+    }
+
+    fn resolve_rate_limit(&self) {
+        let reset_time = *self.rate_limit_count_tx.lock().borrow();
+
+        if let Some(reset_time) = reset_time {
+            if Instant::now() >= reset_time {
+                *self.rate_limit_count_tx.lock().borrow_mut() = None
+            }
+        }
+
+        log::trace!(
+            "resolving reset time: {:?}",
+            *self.rate_limit_count_tx.lock().borrow()
+        );
+    }
+
+    fn update_reset_time(&self, reset_time: Instant) {
+        let original_time = *self.rate_limit_count_tx.lock().borrow();
+
+        let updated_time = if let Some(original_time) = original_time {
+            if reset_time < original_time {
+                Some(reset_time)
+            } else {
+                Some(original_time)
+            }
+        } else {
+            Some(reset_time)
+        };
+
+        log::trace!("updating rate limit time: {:?}", updated_time);
+
+        *self.rate_limit_count_tx.lock().borrow_mut() = updated_time;
+    }
+    async fn send_request(
+        &self,
+        api_key: &str,
+        spans: Vec<&str>,
+        request_timeout: u64,
+    ) -> Result<Response<AsyncBody>> {
+        let request = Request::post("https://api.openai.com/v1/embeddings")
+            .redirect_policy(isahc::config::RedirectPolicy::Follow)
+            .timeout(Duration::from_secs(request_timeout))
+            .header("Content-Type", "application/json")
+            .header("Authorization", format!("Bearer {}", api_key))
+            .body(
+                serde_json::to_string(&OpenAIEmbeddingRequest {
+                    input: spans.clone(),
+                    model: "text-embedding-ada-002",
+                })
+                .unwrap()
+                .into(),
+            )?;
+
+        Ok(self.client.send(request).await?)
+    }
+}
+
+impl CredentialProvider for OpenAIEmbeddingProvider {
+    fn has_credentials(&self) -> bool {
+        match *self.credential.read() {
+            ProviderCredential::Credentials { .. } => true,
+            _ => false,
+        }
+    }
+    fn retrieve_credentials(&self, cx: &AppContext) -> ProviderCredential {
+        let mut credential = self.credential.write();
+        match *credential {
+            ProviderCredential::Credentials { .. } => {
+                return credential.clone();
+            }
+            _ => {
+                if let Ok(api_key) = env::var("OPENAI_API_KEY") {
+                    *credential = ProviderCredential::Credentials { api_key };
+                } else if let Some((_, api_key)) = cx
+                    .platform()
+                    .read_credentials(OPENAI_API_URL)
+                    .log_err()
+                    .flatten()
+                {
+                    if let Some(api_key) = String::from_utf8(api_key).log_err() {
+                        *credential = ProviderCredential::Credentials { api_key };
+                    }
+                } else {
+                };
+            }
+        }
+
+        credential.clone()
+    }
+
+    fn save_credentials(&self, cx: &AppContext, credential: ProviderCredential) {
+        match credential.clone() {
+            ProviderCredential::Credentials { api_key } => {
+                cx.platform()
+                    .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
+                    .log_err();
+            }
+            _ => {}
+        }
+
+        *self.credential.write() = credential;
+    }
+    fn delete_credentials(&self, cx: &AppContext) {
+        cx.platform().delete_credentials(OPENAI_API_URL).log_err();
+        *self.credential.write() = ProviderCredential::NoCredentials;
+    }
+}
+
+#[async_trait]
+impl EmbeddingProvider for OpenAIEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(self.model.clone());
+        model
+    }
+
+    fn max_tokens_per_batch(&self) -> usize {
+        50000
+    }
+
+    fn rate_limit_expiration(&self) -> Option<Instant> {
+        *self.rate_limit_count_rx.borrow()
+    }
+
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
+        const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
+        const MAX_RETRIES: usize = 4;
+
+        let api_key = self.get_api_key()?;
+
+        let mut request_number = 0;
+        let mut rate_limiting = false;
+        let mut request_timeout: u64 = 15;
+        let mut response: Response<AsyncBody>;
+        while request_number < MAX_RETRIES {
+            response = self
+                .send_request(
+                    &api_key,
+                    spans.iter().map(|x| &**x).collect(),
+                    request_timeout,
+                )
+                .await?;
+
+            request_number += 1;
+
+            match response.status() {
+                StatusCode::REQUEST_TIMEOUT => {
+                    request_timeout += 5;
+                }
+                StatusCode::OK => {
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+                    let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
+
+                    log::trace!(
+                        "openai embedding completed. tokens: {:?}",
+                        response.usage.total_tokens
+                    );
+
+                    // If we complete a request successfully that was previously rate_limited
+                    // resolve the rate limit
+                    if rate_limiting {
+                        self.resolve_rate_limit()
+                    }
+
+                    return Ok(response
+                        .data
+                        .into_iter()
+                        .map(|embedding| Embedding::from(embedding.embedding))
+                        .collect());
+                }
+                StatusCode::TOO_MANY_REQUESTS => {
+                    rate_limiting = true;
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+
+                    let delay_duration = {
+                        let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
+                        if let Some(time_to_reset) =
+                            response.headers().get("x-ratelimit-reset-tokens")
+                        {
+                            if let Ok(time_str) = time_to_reset.to_str() {
+                                parse(time_str).unwrap_or(delay)
+                            } else {
+                                delay
+                            }
+                        } else {
+                            delay
+                        }
+                    };
+
+                    // If we've previously rate limited, increment the duration but not the count
+                    let reset_time = Instant::now().add(delay_duration);
+                    self.update_reset_time(reset_time);
+
+                    log::trace!(
+                        "openai rate limiting: waiting {:?} until lifted",
+                        &delay_duration
+                    );
+
+                    self.executor.timer(delay_duration).await;
+                }
+                _ => {
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+                    return Err(anyhow!(
+                        "open ai bad request: {:?} {:?}",
+                        &response.status(),
+                        body
+                    ));
+                }
+            }
+        }
+        Err(anyhow!("openai max retries"))
+    }
+}

crates/ai/src/providers/open_ai/mod.rs

@@ -0,0 +1,9 @@
+pub mod completion;
+pub mod embedding;
+pub mod model;
+
+pub use completion::*;
+pub use embedding::*;
+pub use model::OpenAILanguageModel;
+
+pub const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";

crates/ai/src/providers/open_ai/model.rs

@@ -0,0 +1,57 @@
+use anyhow::anyhow;
+use tiktoken_rs::CoreBPE;
+use util::ResultExt;
+
+use crate::models::{LanguageModel, TruncationDirection};
+
+#[derive(Clone)]
+pub struct OpenAILanguageModel {
+    name: String,
+    bpe: Option<CoreBPE>,
+}
+
+impl OpenAILanguageModel {
+    pub fn load(model_name: &str) -> Self {
+        let bpe = tiktoken_rs::get_bpe_from_model(model_name).log_err();
+        OpenAILanguageModel {
+            name: model_name.to_string(),
+            bpe,
+        }
+    }
+}
+
+impl LanguageModel for OpenAILanguageModel {
+    fn name(&self) -> String {
+        self.name.clone()
+    }
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
+        if let Some(bpe) = &self.bpe {
+            anyhow::Ok(bpe.encode_with_special_tokens(content).len())
+        } else {
+            Err(anyhow!("bpe for open ai model was not retrieved"))
+        }
+    }
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String> {
+        if let Some(bpe) = &self.bpe {
+            let tokens = bpe.encode_with_special_tokens(content);
+            if tokens.len() > length {
+                match direction {
+                    TruncationDirection::End => bpe.decode(tokens[..length].to_vec()),
+                    TruncationDirection::Start => bpe.decode(tokens[length..].to_vec()),
+                }
+            } else {
+                bpe.decode(tokens)
+            }
+        } else {
+            Err(anyhow!("bpe for open ai model was not retrieved"))
+        }
+    }
+    fn capacity(&self) -> anyhow::Result<usize> {
+        anyhow::Ok(tiktoken_rs::model::get_context_size(&self.name))
+    }
+}

crates/ai/src/providers/open_ai/new.rs

@@ -0,0 +1,11 @@
+pub trait LanguageModel {
+    fn name(&self) -> String;
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize>;
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String>;
+    fn capacity(&self) -> anyhow::Result<usize>;
+}

crates/ai/src/test.rs

@@ -0,0 +1,191 @@
+use std::{
+    sync::atomic::{self, AtomicUsize, Ordering},
+    time::Instant,
+};
+
+use async_trait::async_trait;
+use futures::{channel::mpsc, future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
+use gpui::AppContext;
+use parking_lot::Mutex;
+
+use crate::{
+    auth::{CredentialProvider, ProviderCredential},
+    completion::{CompletionProvider, CompletionRequest},
+    embedding::{Embedding, EmbeddingProvider},
+    models::{LanguageModel, TruncationDirection},
+};
+
+#[derive(Clone)]
+pub struct FakeLanguageModel {
+    pub capacity: usize,
+}
+
+impl LanguageModel for FakeLanguageModel {
+    fn name(&self) -> String {
+        "dummy".to_string()
+    }
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
+        anyhow::Ok(content.chars().collect::<Vec<char>>().len())
+    }
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String> {
+        println!("TRYING TO TRUNCATE: {:?}", length.clone());
+
+        if length > self.count_tokens(content)? {
+            println!("NOT TRUNCATING");
+            return anyhow::Ok(content.to_string());
+        }
+
+        anyhow::Ok(match direction {
+            TruncationDirection::End => content.chars().collect::<Vec<char>>()[..length]
+                .into_iter()
+                .collect::<String>(),
+            TruncationDirection::Start => content.chars().collect::<Vec<char>>()[length..]
+                .into_iter()
+                .collect::<String>(),
+        })
+    }
+    fn capacity(&self) -> anyhow::Result<usize> {
+        anyhow::Ok(self.capacity)
+    }
+}
+
+pub struct FakeEmbeddingProvider {
+    pub embedding_count: AtomicUsize,
+}
+
+impl Clone for FakeEmbeddingProvider {
+    fn clone(&self) -> Self {
+        FakeEmbeddingProvider {
+            embedding_count: AtomicUsize::new(self.embedding_count.load(Ordering::SeqCst)),
+        }
+    }
+}
+
+impl Default for FakeEmbeddingProvider {
+    fn default() -> Self {
+        FakeEmbeddingProvider {
+            embedding_count: AtomicUsize::default(),
+        }
+    }
+}
+
+impl FakeEmbeddingProvider {
+    pub fn embedding_count(&self) -> usize {
+        self.embedding_count.load(atomic::Ordering::SeqCst)
+    }
+
+    pub fn embed_sync(&self, span: &str) -> Embedding {
+        let mut result = vec![1.0; 26];
+        for letter in span.chars() {
+            let letter = letter.to_ascii_lowercase();
+            if letter as u32 >= 'a' as u32 {
+                let ix = (letter as u32) - ('a' as u32);
+                if ix < 26 {
+                    result[ix as usize] += 1.0;
+                }
+            }
+        }
+
+        let norm = result.iter().map(|x| x * x).sum::<f32>().sqrt();
+        for x in &mut result {
+            *x /= norm;
+        }
+
+        result.into()
+    }
+}
+
+impl CredentialProvider for FakeEmbeddingProvider {
+    fn has_credentials(&self) -> bool {
+        true
+    }
+    fn retrieve_credentials(&self, _cx: &AppContext) -> ProviderCredential {
+        ProviderCredential::NotNeeded
+    }
+    fn save_credentials(&self, _cx: &AppContext, _credential: ProviderCredential) {}
+    fn delete_credentials(&self, _cx: &AppContext) {}
+}
+
+#[async_trait]
+impl EmbeddingProvider for FakeEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        Box::new(FakeLanguageModel { capacity: 1000 })
+    }
+    fn max_tokens_per_batch(&self) -> usize {
+        1000
+    }
+
+    fn rate_limit_expiration(&self) -> Option<Instant> {
+        None
+    }
+
+    async fn embed_batch(&self, spans: Vec<String>) -> anyhow::Result<Vec<Embedding>> {
+        self.embedding_count
+            .fetch_add(spans.len(), atomic::Ordering::SeqCst);
+
+        anyhow::Ok(spans.iter().map(|span| self.embed_sync(span)).collect())
+    }
+}
+
+pub struct FakeCompletionProvider {
+    last_completion_tx: Mutex<Option<mpsc::Sender<String>>>,
+}
+
+impl Clone for FakeCompletionProvider {
+    fn clone(&self) -> Self {
+        Self {
+            last_completion_tx: Mutex::new(None),
+        }
+    }
+}
+
+impl FakeCompletionProvider {
+    pub fn new() -> Self {
+        Self {
+            last_completion_tx: Mutex::new(None),
+        }
+    }
+
+    pub fn send_completion(&self, completion: impl Into<String>) {
+        let mut tx = self.last_completion_tx.lock();
+        tx.as_mut().unwrap().try_send(completion.into()).unwrap();
+    }
+
+    pub fn finish_completion(&self) {
+        self.last_completion_tx.lock().take().unwrap();
+    }
+}
+
+impl CredentialProvider for FakeCompletionProvider {
+    fn has_credentials(&self) -> bool {
+        true
+    }
+    fn retrieve_credentials(&self, _cx: &AppContext) -> ProviderCredential {
+        ProviderCredential::NotNeeded
+    }
+    fn save_credentials(&self, _cx: &AppContext, _credential: ProviderCredential) {}
+    fn delete_credentials(&self, _cx: &AppContext) {}
+}
+
+impl CompletionProvider for FakeCompletionProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(FakeLanguageModel { capacity: 8190 });
+        model
+    }
+    fn complete(
+        &self,
+        _prompt: Box<dyn CompletionRequest>,
+    ) -> BoxFuture<'static, anyhow::Result<BoxStream<'static, anyhow::Result<String>>>> {
+        let (tx, rx) = mpsc::channel(1);
+        *self.last_completion_tx.lock() = Some(tx);
+        async move { Ok(rx.map(|rx| Ok(rx)).boxed()) }.boxed()
+    }
+    fn box_clone(&self) -> Box<dyn CompletionProvider> {
+        Box::new((*self).clone())
+    }
+}

crates/ai2/Cargo.toml

@@ -0,0 +1,38 @@
+[package]
+name = "ai2"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[lib]
+path = "src/ai2.rs"
+doctest = false
+
+[features]
+test-support = []
+
+[dependencies]
+gpui2 = { path = "../gpui2" }
+util = { path = "../util" }
+language2 = { path = "../language2" }
+async-trait.workspace = true
+anyhow.workspace = true
+futures.workspace = true
+lazy_static.workspace = true
+ordered-float.workspace = true
+parking_lot.workspace = true
+isahc.workspace = true
+regex.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+postage.workspace = true
+rand.workspace = true
+log.workspace = true
+parse_duration = "2.1.1"
+tiktoken-rs = "0.5.0"
+matrixmultiply = "0.3.7"
+rusqlite = { version = "0.29.0", features = ["blob", "array", "modern_sqlite"] }
+bincode = "1.3.3"
+
+[dev-dependencies]
+gpui2 = { path = "../gpui2", features = ["test-support"] }

crates/ai2/src/ai2.rs

@@ -0,0 +1,8 @@
+pub mod auth;
+pub mod completion;
+pub mod embedding;
+pub mod models;
+pub mod prompts;
+pub mod providers;
+#[cfg(any(test, feature = "test-support"))]
+pub mod test;

crates/ai2/src/auth.rs

@@ -0,0 +1,17 @@
+use async_trait::async_trait;
+use gpui2::AppContext;
+
+#[derive(Clone, Debug)]
+pub enum ProviderCredential {
+    Credentials { api_key: String },
+    NoCredentials,
+    NotNeeded,
+}
+
+#[async_trait]
+pub trait CredentialProvider: Send + Sync {
+    fn has_credentials(&self) -> bool;
+    async fn retrieve_credentials(&self, cx: &mut AppContext) -> ProviderCredential;
+    async fn save_credentials(&self, cx: &mut AppContext, credential: ProviderCredential);
+    async fn delete_credentials(&self, cx: &mut AppContext);
+}

crates/ai2/src/completion.rs

@@ -0,0 +1,23 @@
+use anyhow::Result;
+use futures::{future::BoxFuture, stream::BoxStream};
+
+use crate::{auth::CredentialProvider, models::LanguageModel};
+
+pub trait CompletionRequest: Send + Sync {
+    fn data(&self) -> serde_json::Result<String>;
+}
+
+pub trait CompletionProvider: CredentialProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel>;
+    fn complete(
+        &self,
+        prompt: Box<dyn CompletionRequest>,
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
+    fn box_clone(&self) -> Box<dyn CompletionProvider>;
+}
+
+impl Clone for Box<dyn CompletionProvider> {
+    fn clone(&self) -> Box<dyn CompletionProvider> {
+        self.box_clone()
+    }
+}

crates/ai2/src/embedding.rs

@@ -0,0 +1,123 @@
+use std::time::Instant;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use ordered_float::OrderedFloat;
+use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef};
+use rusqlite::ToSql;
+
+use crate::auth::CredentialProvider;
+use crate::models::LanguageModel;
+
+#[derive(Debug, PartialEq, Clone)]
+pub struct Embedding(pub Vec<f32>);
+
+// This is needed for semantic index functionality
+// Unfortunately it has to live wherever the "Embedding" struct is created.
+// Keeping this in here though, introduces a 'rusqlite' dependency into AI
+// which is less than ideal
+impl FromSql for Embedding {
+    fn column_result(value: ValueRef) -> FromSqlResult<Self> {
+        let bytes = value.as_blob()?;
+        let embedding: Result<Vec<f32>, Box<bincode::ErrorKind>> = bincode::deserialize(bytes);
+        if embedding.is_err() {
+            return Err(rusqlite::types::FromSqlError::Other(embedding.unwrap_err()));
+        }
+        Ok(Embedding(embedding.unwrap()))
+    }
+}
+
+impl ToSql for Embedding {
+    fn to_sql(&self) -> rusqlite::Result<ToSqlOutput> {
+        let bytes = bincode::serialize(&self.0)
+            .map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
+        Ok(ToSqlOutput::Owned(rusqlite::types::Value::Blob(bytes)))
+    }
+}
+impl From<Vec<f32>> for Embedding {
+    fn from(value: Vec<f32>) -> Self {
+        Embedding(value)
+    }
+}
+
+impl Embedding {
+    pub fn similarity(&self, other: &Self) -> OrderedFloat<f32> {
+        let len = self.0.len();
+        assert_eq!(len, other.0.len());
+
+        let mut result = 0.0;
+        unsafe {
+            matrixmultiply::sgemm(
+                1,
+                len,
+                1,
+                1.0,
+                self.0.as_ptr(),
+                len as isize,
+                1,
+                other.0.as_ptr(),
+                1,
+                len as isize,
+                0.0,
+                &mut result as *mut f32,
+                1,
+                1,
+            );
+        }
+        OrderedFloat(result)
+    }
+}
+
+#[async_trait]
+pub trait EmbeddingProvider: CredentialProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel>;
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
+    fn max_tokens_per_batch(&self) -> usize;
+    fn rate_limit_expiration(&self) -> Option<Instant>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::prelude::*;
+
+    #[gpui2::test]
+    fn test_similarity(mut rng: StdRng) {
+        assert_eq!(
+            Embedding::from(vec![1., 0., 0., 0., 0.])
+                .similarity(&Embedding::from(vec![0., 1., 0., 0., 0.])),
+            0.
+        );
+        assert_eq!(
+            Embedding::from(vec![2., 0., 0., 0., 0.])
+                .similarity(&Embedding::from(vec![3., 1., 0., 0., 0.])),
+            6.
+        );
+
+        for _ in 0..100 {
+            let size = 1536;
+            let mut a = vec![0.; size];
+            let mut b = vec![0.; size];
+            for (a, b) in a.iter_mut().zip(b.iter_mut()) {
+                *a = rng.gen();
+                *b = rng.gen();
+            }
+            let a = Embedding::from(a);
+            let b = Embedding::from(b);
+
+            assert_eq!(
+                round_to_decimals(a.similarity(&b), 1),
+                round_to_decimals(reference_dot(&a.0, &b.0), 1)
+            );
+        }
+
+        fn round_to_decimals(n: OrderedFloat<f32>, decimal_places: i32) -> f32 {
+            let factor = (10.0 as f32).powi(decimal_places);
+            (n * factor).round() / factor
+        }
+
+        fn reference_dot(a: &[f32], b: &[f32]) -> OrderedFloat<f32> {
+            OrderedFloat(a.iter().zip(b.iter()).map(|(a, b)| a * b).sum())
+        }
+    }
+}

crates/ai2/src/models.rs

@@ -0,0 +1,16 @@
+pub enum TruncationDirection {
+    Start,
+    End,
+}
+
+pub trait LanguageModel {
+    fn name(&self) -> String;
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize>;
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String>;
+    fn capacity(&self) -> anyhow::Result<usize>;
+}

crates/ai2/src/prompts/base.rs

@@ -0,0 +1,330 @@
+use std::cmp::Reverse;
+use std::ops::Range;
+use std::sync::Arc;
+
+use language2::BufferSnapshot;
+use util::ResultExt;
+
+use crate::models::LanguageModel;
+use crate::prompts::repository_context::PromptCodeSnippet;
+
+pub(crate) enum PromptFileType {
+    Text,
+    Code,
+}
+
+// TODO: Set this up to manage for defaults well
+pub struct PromptArguments {
+    pub model: Arc<dyn LanguageModel>,
+    pub user_prompt: Option<String>,
+    pub language_name: Option<String>,
+    pub project_name: Option<String>,
+    pub snippets: Vec<PromptCodeSnippet>,
+    pub reserved_tokens: usize,
+    pub buffer: Option<BufferSnapshot>,
+    pub selected_range: Option<Range<usize>>,
+}
+
+impl PromptArguments {
+    pub(crate) fn get_file_type(&self) -> PromptFileType {
+        if self
+            .language_name
+            .as_ref()
+            .and_then(|name| Some(!["Markdown", "Plain Text"].contains(&name.as_str())))
+            .unwrap_or(true)
+        {
+            PromptFileType::Code
+        } else {
+            PromptFileType::Text
+        }
+    }
+}
+
+pub trait PromptTemplate {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)>;
+}
+
+#[repr(i8)]
+#[derive(PartialEq, Eq, Ord)]
+pub enum PromptPriority {
+    Mandatory,                // Ignores truncation
+    Ordered { order: usize }, // Truncates based on priority
+}
+
+impl PartialOrd for PromptPriority {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        match (self, other) {
+            (Self::Mandatory, Self::Mandatory) => Some(std::cmp::Ordering::Equal),
+            (Self::Mandatory, Self::Ordered { .. }) => Some(std::cmp::Ordering::Greater),
+            (Self::Ordered { .. }, Self::Mandatory) => Some(std::cmp::Ordering::Less),
+            (Self::Ordered { order: a }, Self::Ordered { order: b }) => b.partial_cmp(a),
+        }
+    }
+}
+
+pub struct PromptChain {
+    args: PromptArguments,
+    templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
+}
+
+impl PromptChain {
+    pub fn new(
+        args: PromptArguments,
+        templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
+    ) -> Self {
+        PromptChain { args, templates }
+    }
+
+    pub fn generate(&self, truncate: bool) -> anyhow::Result<(String, usize)> {
+        // Argsort based on Prompt Priority
+        let seperator = "\n";
+        let seperator_tokens = self.args.model.count_tokens(seperator)?;
+        let mut sorted_indices = (0..self.templates.len()).collect::<Vec<_>>();
+        sorted_indices.sort_by_key(|&i| Reverse(&self.templates[i].0));
+
+        // If Truncate
+        let mut tokens_outstanding = if truncate {
+            Some(self.args.model.capacity()? - self.args.reserved_tokens)
+        } else {
+            None
+        };
+
+        let mut prompts = vec!["".to_string(); sorted_indices.len()];
+        for idx in sorted_indices {
+            let (_, template) = &self.templates[idx];
+
+            if let Some((template_prompt, prompt_token_count)) =
+                template.generate(&self.args, tokens_outstanding).log_err()
+            {
+                if template_prompt != "" {
+                    prompts[idx] = template_prompt;
+
+                    if let Some(remaining_tokens) = tokens_outstanding {
+                        let new_tokens = prompt_token_count + seperator_tokens;
+                        tokens_outstanding = if remaining_tokens > new_tokens {
+                            Some(remaining_tokens - new_tokens)
+                        } else {
+                            Some(0)
+                        };
+                    }
+                }
+            }
+        }
+
+        prompts.retain(|x| x != "");
+
+        let full_prompt = prompts.join(seperator);
+        let total_token_count = self.args.model.count_tokens(&full_prompt)?;
+        anyhow::Ok((prompts.join(seperator), total_token_count))
+    }
+}
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use crate::models::TruncationDirection;
+    use crate::test::FakeLanguageModel;
+
+    use super::*;
+
+    #[test]
+    pub fn test_prompt_chain() {
+        struct TestPromptTemplate {}
+        impl PromptTemplate for TestPromptTemplate {
+            fn generate(
+                &self,
+                args: &PromptArguments,
+                max_token_length: Option<usize>,
+            ) -> anyhow::Result<(String, usize)> {
+                let mut content = "This is a test prompt template".to_string();
+
+                let mut token_count = args.model.count_tokens(&content)?;
+                if let Some(max_token_length) = max_token_length {
+                    if token_count > max_token_length {
+                        content = args.model.truncate(
+                            &content,
+                            max_token_length,
+                            TruncationDirection::End,
+                        )?;
+                        token_count = max_token_length;
+                    }
+                }
+
+                anyhow::Ok((content, token_count))
+            }
+        }
+
+        struct TestLowPriorityTemplate {}
+        impl PromptTemplate for TestLowPriorityTemplate {
+            fn generate(
+                &self,
+                args: &PromptArguments,
+                max_token_length: Option<usize>,
+            ) -> anyhow::Result<(String, usize)> {
+                let mut content = "This is a low priority test prompt template".to_string();
+
+                let mut token_count = args.model.count_tokens(&content)?;
+                if let Some(max_token_length) = max_token_length {
+                    if token_count > max_token_length {
+                        content = args.model.truncate(
+                            &content,
+                            max_token_length,
+                            TruncationDirection::End,
+                        )?;
+                        token_count = max_token_length;
+                    }
+                }
+
+                anyhow::Ok((content, token_count))
+            }
+        }
+
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity: 100 });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(false).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a test prompt template\nThis is a low priority test prompt template"
+                .to_string()
+        );
+
+        assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
+
+        // Testing with Truncation Off
+        // Should ignore capacity and return all prompts
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity: 20 });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(false).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a test prompt template\nThis is a low priority test prompt template"
+                .to_string()
+        );
+
+        assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
+
+        // Testing with Truncation Off
+        // Should ignore capacity and return all prompts
+        let capacity = 20;
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 2 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(true).unwrap();
+
+        assert_eq!(prompt, "This is a test promp".to_string());
+        assert_eq!(token_count, capacity);
+
+        // Change Ordering of Prompts Based on Priority
+        let capacity = 120;
+        let reserved_tokens = 10;
+        let model: Arc<dyn LanguageModel> = Arc::new(FakeLanguageModel { capacity });
+        let args = PromptArguments {
+            model: model.clone(),
+            language_name: None,
+            project_name: None,
+            snippets: Vec::new(),
+            reserved_tokens,
+            buffer: None,
+            selected_range: None,
+            user_prompt: None,
+        };
+        let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+            (
+                PromptPriority::Mandatory,
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 0 },
+                Box::new(TestPromptTemplate {}),
+            ),
+            (
+                PromptPriority::Ordered { order: 1 },
+                Box::new(TestLowPriorityTemplate {}),
+            ),
+        ];
+        let chain = PromptChain::new(args, templates);
+
+        let (prompt, token_count) = chain.generate(true).unwrap();
+
+        assert_eq!(
+            prompt,
+            "This is a low priority test prompt template\nThis is a test prompt template\nThis is a low priority test prompt "
+                .to_string()
+        );
+        assert_eq!(token_count, capacity - reserved_tokens);
+    }
+}

crates/ai2/src/prompts/file_context.rs

@@ -0,0 +1,164 @@
+use anyhow::anyhow;
+use language2::BufferSnapshot;
+use language2::ToOffset;
+
+use crate::models::LanguageModel;
+use crate::models::TruncationDirection;
+use crate::prompts::base::PromptArguments;
+use crate::prompts::base::PromptTemplate;
+use std::fmt::Write;
+use std::ops::Range;
+use std::sync::Arc;
+
+fn retrieve_context(
+    buffer: &BufferSnapshot,
+    selected_range: &Option<Range<usize>>,
+    model: Arc<dyn LanguageModel>,
+    max_token_count: Option<usize>,
+) -> anyhow::Result<(String, usize, bool)> {
+    let mut prompt = String::new();
+    let mut truncated = false;
+    if let Some(selected_range) = selected_range {
+        let start = selected_range.start.to_offset(buffer);
+        let end = selected_range.end.to_offset(buffer);
+
+        let start_window = buffer.text_for_range(0..start).collect::<String>();
+
+        let mut selected_window = String::new();
+        if start == end {
+            write!(selected_window, "<|START|>").unwrap();
+        } else {
+            write!(selected_window, "<|START|").unwrap();
+        }
+
+        write!(
+            selected_window,
+            "{}",
+            buffer.text_for_range(start..end).collect::<String>()
+        )
+        .unwrap();
+
+        if start != end {
+            write!(selected_window, "|END|>").unwrap();
+        }
+
+        let end_window = buffer.text_for_range(end..buffer.len()).collect::<String>();
+
+        if let Some(max_token_count) = max_token_count {
+            let selected_tokens = model.count_tokens(&selected_window)?;
+            if selected_tokens > max_token_count {
+                return Err(anyhow!(
+                    "selected range is greater than model context window, truncation not possible"
+                ));
+            };
+
+            let mut remaining_tokens = max_token_count - selected_tokens;
+            let start_window_tokens = model.count_tokens(&start_window)?;
+            let end_window_tokens = model.count_tokens(&end_window)?;
+            let outside_tokens = start_window_tokens + end_window_tokens;
+            if outside_tokens > remaining_tokens {
+                let (start_goal_tokens, end_goal_tokens) =
+                    if start_window_tokens < end_window_tokens {
+                        let start_goal_tokens = (remaining_tokens / 2).min(start_window_tokens);
+                        remaining_tokens -= start_goal_tokens;
+                        let end_goal_tokens = remaining_tokens.min(end_window_tokens);
+                        (start_goal_tokens, end_goal_tokens)
+                    } else {
+                        let end_goal_tokens = (remaining_tokens / 2).min(end_window_tokens);
+                        remaining_tokens -= end_goal_tokens;
+                        let start_goal_tokens = remaining_tokens.min(start_window_tokens);
+                        (start_goal_tokens, end_goal_tokens)
+                    };
+
+                let truncated_start_window =
+                    model.truncate(&start_window, start_goal_tokens, TruncationDirection::Start)?;
+                let truncated_end_window =
+                    model.truncate(&end_window, end_goal_tokens, TruncationDirection::End)?;
+                writeln!(
+                    prompt,
+                    "{truncated_start_window}{selected_window}{truncated_end_window}"
+                )
+                .unwrap();
+                truncated = true;
+            } else {
+                writeln!(prompt, "{start_window}{selected_window}{end_window}").unwrap();
+            }
+        } else {
+            // If we dont have a selected range, include entire file.
+            writeln!(prompt, "{}", &buffer.text()).unwrap();
+
+            // Dumb truncation strategy
+            if let Some(max_token_count) = max_token_count {
+                if model.count_tokens(&prompt)? > max_token_count {
+                    truncated = true;
+                    prompt = model.truncate(&prompt, max_token_count, TruncationDirection::End)?;
+                }
+            }
+        }
+    }
+
+    let token_count = model.count_tokens(&prompt)?;
+    anyhow::Ok((prompt, token_count, truncated))
+}
+
+pub struct FileContext {}
+
+impl PromptTemplate for FileContext {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        if let Some(buffer) = &args.buffer {
+            let mut prompt = String::new();
+            // Add Initial Preamble
+            // TODO: Do we want to add the path in here?
+            writeln!(
+                prompt,
+                "The file you are currently working on has the following content:"
+            )
+            .unwrap();
+
+            let language_name = args
+                .language_name
+                .clone()
+                .unwrap_or("".to_string())
+                .to_lowercase();
+
+            let (context, _, truncated) = retrieve_context(
+                buffer,
+                &args.selected_range,
+                args.model.clone(),
+                max_token_length,
+            )?;
+            writeln!(prompt, "```{language_name}\n{context}\n```").unwrap();
+
+            if truncated {
+                writeln!(prompt, "Note the content has been truncated and only represents a portion of the file.").unwrap();
+            }
+
+            if let Some(selected_range) = &args.selected_range {
+                let start = selected_range.start.to_offset(buffer);
+                let end = selected_range.end.to_offset(buffer);
+
+                if start == end {
+                    writeln!(prompt, "In particular, the user's cursor is currently on the '<|START|>' span in the above content, with no text selected.").unwrap();
+                } else {
+                    writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
+                }
+            }
+
+            // Really dumb truncation strategy
+            if let Some(max_tokens) = max_token_length {
+                prompt = args
+                    .model
+                    .truncate(&prompt, max_tokens, TruncationDirection::End)?;
+            }
+
+            let token_count = args.model.count_tokens(&prompt)?;
+            anyhow::Ok((prompt, token_count))
+        } else {
+            Err(anyhow!("no buffer provided to retrieve file context from"))
+        }
+    }
+}

crates/ai2/src/prompts/generate.rs

@@ -0,0 +1,99 @@
+use crate::prompts::base::{PromptArguments, PromptFileType, PromptTemplate};
+use anyhow::anyhow;
+use std::fmt::Write;
+
+pub fn capitalize(s: &str) -> String {
+    let mut c = s.chars();
+    match c.next() {
+        None => String::new(),
+        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
+    }
+}
+
+pub struct GenerateInlineContent {}
+
+impl PromptTemplate for GenerateInlineContent {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let Some(user_prompt) = &args.user_prompt else {
+            return Err(anyhow!("user prompt not provided"));
+        };
+
+        let file_type = args.get_file_type();
+        let content_type = match &file_type {
+            PromptFileType::Code => "code",
+            PromptFileType::Text => "text",
+        };
+
+        let mut prompt = String::new();
+
+        if let Some(selected_range) = &args.selected_range {
+            if selected_range.start == selected_range.end {
+                writeln!(
+                    prompt,
+                    "Assume the cursor is located where the `<|START|>` span is."
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "{} can't be replaced, so assume your answer will be inserted at the cursor.",
+                    capitalize(content_type)
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "Generate {content_type} based on the users prompt: {user_prompt}",
+                )
+                .unwrap();
+            } else {
+                writeln!(prompt, "Modify the user's selected {content_type} based upon the users prompt: '{user_prompt}'").unwrap();
+                writeln!(prompt, "You must reply with only the adjusted {content_type} (within the '<|START|' and '|END|>' spans) not the entire file.").unwrap();
+                writeln!(prompt, "Double check that you only return code and not the '<|START|' and '|END|'> spans").unwrap();
+            }
+        } else {
+            writeln!(
+                prompt,
+                "Generate {content_type} based on the users prompt: {user_prompt}"
+            )
+            .unwrap();
+        }
+
+        if let Some(language_name) = &args.language_name {
+            writeln!(
+                prompt,
+                "Your answer MUST always and only be valid {}.",
+                language_name
+            )
+            .unwrap();
+        }
+        writeln!(prompt, "Never make remarks about the output.").unwrap();
+        writeln!(
+            prompt,
+            "Do not return anything else, except the generated {content_type}."
+        )
+        .unwrap();
+
+        match file_type {
+            PromptFileType::Code => {
+                // writeln!(prompt, "Always wrap your code in a Markdown block.").unwrap();
+            }
+            _ => {}
+        }
+
+        // Really dumb truncation strategy
+        if let Some(max_tokens) = max_token_length {
+            prompt = args.model.truncate(
+                &prompt,
+                max_tokens,
+                crate::models::TruncationDirection::End,
+            )?;
+        }
+
+        let token_count = args.model.count_tokens(&prompt)?;
+
+        anyhow::Ok((prompt, token_count))
+    }
+}

crates/ai2/src/prompts/mod.rs

@@ -0,0 +1,5 @@
+pub mod base;
+pub mod file_context;
+pub mod generate;
+pub mod preamble;
+pub mod repository_context;

crates/ai2/src/prompts/preamble.rs

@@ -0,0 +1,52 @@
+use crate::prompts::base::{PromptArguments, PromptFileType, PromptTemplate};
+use std::fmt::Write;
+
+pub struct EngineerPreamble {}
+
+impl PromptTemplate for EngineerPreamble {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let mut prompts = Vec::new();
+
+        match args.get_file_type() {
+            PromptFileType::Code => {
+                prompts.push(format!(
+                    "You are an expert {}engineer.",
+                    args.language_name.clone().unwrap_or("".to_string()) + " "
+                ));
+            }
+            PromptFileType::Text => {
+                prompts.push("You are an expert engineer.".to_string());
+            }
+        }
+
+        if let Some(project_name) = args.project_name.clone() {
+            prompts.push(format!(
+                "You are currently working inside the '{project_name}' project in code editor Zed."
+            ));
+        }
+
+        if let Some(mut remaining_tokens) = max_token_length {
+            let mut prompt = String::new();
+            let mut total_count = 0;
+            for prompt_piece in prompts {
+                let prompt_token_count =
+                    args.model.count_tokens(&prompt_piece)? + args.model.count_tokens("\n")?;
+                if remaining_tokens > prompt_token_count {
+                    writeln!(prompt, "{prompt_piece}").unwrap();
+                    remaining_tokens -= prompt_token_count;
+                    total_count += prompt_token_count;
+                }
+            }
+
+            anyhow::Ok((prompt, total_count))
+        } else {
+            let prompt = prompts.join("\n");
+            let token_count = args.model.count_tokens(&prompt)?;
+            anyhow::Ok((prompt, token_count))
+        }
+    }
+}

crates/ai2/src/prompts/repository_context.rs

@@ -0,0 +1,98 @@
+use crate::prompts::base::{PromptArguments, PromptTemplate};
+use std::fmt::Write;
+use std::{ops::Range, path::PathBuf};
+
+use gpui2::{AsyncAppContext, Model};
+use language2::{Anchor, Buffer};
+
+#[derive(Clone)]
+pub struct PromptCodeSnippet {
+    path: Option<PathBuf>,
+    language_name: Option<String>,
+    content: String,
+}
+
+impl PromptCodeSnippet {
+    pub fn new(
+        buffer: Model<Buffer>,
+        range: Range<Anchor>,
+        cx: &mut AsyncAppContext,
+    ) -> anyhow::Result<Self> {
+        let (content, language_name, file_path) = buffer.update(cx, |buffer, _| {
+            let snapshot = buffer.snapshot();
+            let content = snapshot.text_for_range(range.clone()).collect::<String>();
+
+            let language_name = buffer
+                .language()
+                .and_then(|language| Some(language.name().to_string().to_lowercase()));
+
+            let file_path = buffer
+                .file()
+                .and_then(|file| Some(file.path().to_path_buf()));
+
+            (content, language_name, file_path)
+        })?;
+
+        anyhow::Ok(PromptCodeSnippet {
+            path: file_path,
+            language_name,
+            content,
+        })
+    }
+}
+
+impl ToString for PromptCodeSnippet {
+    fn to_string(&self) -> String {
+        let path = self
+            .path
+            .as_ref()
+            .and_then(|path| Some(path.to_string_lossy().to_string()))
+            .unwrap_or("".to_string());
+        let language_name = self.language_name.clone().unwrap_or("".to_string());
+        let content = self.content.clone();
+
+        format!("The below code snippet may be relevant from file: {path}\n```{language_name}\n{content}\n```")
+    }
+}
+
+pub struct RepositoryContext {}
+
+impl PromptTemplate for RepositoryContext {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        const MAXIMUM_SNIPPET_TOKEN_COUNT: usize = 500;
+        let template = "You are working inside a large repository, here are a few code snippets that may be useful.";
+        let mut prompt = String::new();
+
+        let mut remaining_tokens = max_token_length.clone();
+        let seperator_token_length = args.model.count_tokens("\n")?;
+        for snippet in &args.snippets {
+            let mut snippet_prompt = template.to_string();
+            let content = snippet.to_string();
+            writeln!(snippet_prompt, "{content}").unwrap();
+
+            let token_count = args.model.count_tokens(&snippet_prompt)?;
+            if token_count <= MAXIMUM_SNIPPET_TOKEN_COUNT {
+                if let Some(tokens_left) = remaining_tokens {
+                    if tokens_left >= token_count {
+                        writeln!(prompt, "{snippet_prompt}").unwrap();
+                        remaining_tokens = if tokens_left >= (token_count + seperator_token_length)
+                        {
+                            Some(tokens_left - token_count - seperator_token_length)
+                        } else {
+                            Some(0)
+                        };
+                    }
+                } else {
+                    writeln!(prompt, "{snippet_prompt}").unwrap();
+                }
+            }
+        }
+
+        let total_token_count = args.model.count_tokens(&prompt)?;
+        anyhow::Ok((prompt, total_token_count))
+    }
+}

crates/ai2/src/providers/mod.rs

@@ -0,0 +1 @@
+pub mod open_ai;

crates/ai2/src/providers/open_ai/completion.rs

@@ -0,0 +1,306 @@
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use futures::{
+    future::BoxFuture, io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, FutureExt,
+    Stream, StreamExt,
+};
+use gpui2::{AppContext, Executor};
+use isahc::{http::StatusCode, Request, RequestExt};
+use parking_lot::RwLock;
+use serde::{Deserialize, Serialize};
+use std::{
+    env,
+    fmt::{self, Display},
+    io,
+    sync::Arc,
+};
+use util::ResultExt;
+
+use crate::{
+    auth::{CredentialProvider, ProviderCredential},
+    completion::{CompletionProvider, CompletionRequest},
+    models::LanguageModel,
+};
+
+use crate::providers::open_ai::{OpenAILanguageModel, OPENAI_API_URL};
+
+#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    User,
+    Assistant,
+    System,
+}
+
+impl Role {
+    pub fn cycle(&mut self) {
+        *self = match self {
+            Role::User => Role::Assistant,
+            Role::Assistant => Role::System,
+            Role::System => Role::User,
+        }
+    }
+}
+
+impl Display for Role {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Role::User => write!(f, "User"),
+            Role::Assistant => write!(f, "Assistant"),
+            Role::System => write!(f, "System"),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+pub struct RequestMessage {
+    pub role: Role,
+    pub content: String,
+}
+
+#[derive(Debug, Default, Serialize)]
+pub struct OpenAIRequest {
+    pub model: String,
+    pub messages: Vec<RequestMessage>,
+    pub stream: bool,
+    pub stop: Vec<String>,
+    pub temperature: f32,
+}
+
+impl CompletionRequest for OpenAIRequest {
+    fn data(&self) -> serde_json::Result<String> {
+        serde_json::to_string(self)
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+pub struct ResponseMessage {
+    pub role: Option<Role>,
+    pub content: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct OpenAIUsage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct ChatChoiceDelta {
+    pub index: u32,
+    pub delta: ResponseMessage,
+    pub finish_reason: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct OpenAIResponseStreamEvent {
+    pub id: Option<String>,
+    pub object: String,
+    pub created: u32,
+    pub model: String,
+    pub choices: Vec<ChatChoiceDelta>,
+    pub usage: Option<OpenAIUsage>,
+}
+
+pub async fn stream_completion(
+    credential: ProviderCredential,
+    executor: Arc<Executor>,
+    request: Box<dyn CompletionRequest>,
+) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
+    let api_key = match credential {
+        ProviderCredential::Credentials { api_key } => api_key,
+        _ => {
+            return Err(anyhow!("no credentials provider for completion"));
+        }
+    };
+
+    let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
+
+    let json_data = request.data()?;
+    let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
+        .header("Content-Type", "application/json")
+        .header("Authorization", format!("Bearer {}", api_key))
+        .body(json_data)?
+        .send_async()
+        .await?;
+
+    let status = response.status();
+    if status == StatusCode::OK {
+        executor
+            .spawn(async move {
+                let mut lines = BufReader::new(response.body_mut()).lines();
+
+                fn parse_line(
+                    line: Result<String, io::Error>,
+                ) -> Result<Option<OpenAIResponseStreamEvent>> {
+                    if let Some(data) = line?.strip_prefix("data: ") {
+                        let event = serde_json::from_str(&data)?;
+                        Ok(Some(event))
+                    } else {
+                        Ok(None)
+                    }
+                }
+
+                while let Some(line) = lines.next().await {
+                    if let Some(event) = parse_line(line).transpose() {
+                        let done = event.as_ref().map_or(false, |event| {
+                            event
+                                .choices
+                                .last()
+                                .map_or(false, |choice| choice.finish_reason.is_some())
+                        });
+                        if tx.unbounded_send(event).is_err() {
+                            break;
+                        }
+
+                        if done {
+                            break;
+                        }
+                    }
+                }
+
+                anyhow::Ok(())
+            })
+            .detach();
+
+        Ok(rx)
+    } else {
+        let mut body = String::new();
+        response.body_mut().read_to_string(&mut body).await?;
+
+        #[derive(Deserialize)]
+        struct OpenAIResponse {
+            error: OpenAIError,
+        }
+
+        #[derive(Deserialize)]
+        struct OpenAIError {
+            message: String,
+        }
+
+        match serde_json::from_str::<OpenAIResponse>(&body) {
+            Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
+                "Failed to connect to OpenAI API: {}",
+                response.error.message,
+            )),
+
+            _ => Err(anyhow!(
+                "Failed to connect to OpenAI API: {} {}",
+                response.status(),
+                body,
+            )),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct OpenAICompletionProvider {
+    model: OpenAILanguageModel,
+    credential: Arc<RwLock<ProviderCredential>>,
+    executor: Arc<Executor>,
+}
+
+impl OpenAICompletionProvider {
+    pub fn new(model_name: &str, executor: Arc<Executor>) -> Self {
+        let model = OpenAILanguageModel::load(model_name);
+        let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
+        Self {
+            model,
+            credential,
+            executor,
+        }
+    }
+}
+
+#[async_trait]
+impl CredentialProvider for OpenAICompletionProvider {
+    fn has_credentials(&self) -> bool {
+        match *self.credential.read() {
+            ProviderCredential::Credentials { .. } => true,
+            _ => false,
+        }
+    }
+    async fn retrieve_credentials(&self, cx: &mut AppContext) -> ProviderCredential {
+        let existing_credential = self.credential.read().clone();
+
+        let retrieved_credential = cx
+            .run_on_main(move |cx| match existing_credential {
+                ProviderCredential::Credentials { .. } => {
+                    return existing_credential.clone();
+                }
+                _ => {
+                    if let Some(api_key) = env::var("OPENAI_API_KEY").log_err() {
+                        return ProviderCredential::Credentials { api_key };
+                    }
+
+                    if let Some(Some((_, api_key))) = cx.read_credentials(OPENAI_API_URL).log_err()
+                    {
+                        if let Some(api_key) = String::from_utf8(api_key).log_err() {
+                            return ProviderCredential::Credentials { api_key };
+                        } else {
+                            return ProviderCredential::NoCredentials;
+                        }
+                    } else {
+                        return ProviderCredential::NoCredentials;
+                    }
+                }
+            })
+            .await;
+
+        *self.credential.write() = retrieved_credential.clone();
+        retrieved_credential
+    }
+
+    async fn save_credentials(&self, cx: &mut AppContext, credential: ProviderCredential) {
+        *self.credential.write() = credential.clone();
+        let credential = credential.clone();
+        cx.run_on_main(move |cx| match credential {
+            ProviderCredential::Credentials { api_key } => {
+                cx.write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
+                    .log_err();
+            }
+            _ => {}
+        })
+        .await;
+    }
+    async fn delete_credentials(&self, cx: &mut AppContext) {
+        cx.run_on_main(move |cx| cx.delete_credentials(OPENAI_API_URL).log_err())
+            .await;
+        *self.credential.write() = ProviderCredential::NoCredentials;
+    }
+}
+
+impl CompletionProvider for OpenAICompletionProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(self.model.clone());
+        model
+    }
+    fn complete(
+        &self,
+        prompt: Box<dyn CompletionRequest>,
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
+        // Currently the CompletionRequest for OpenAI, includes a 'model' parameter
+        // This means that the model is determined by the CompletionRequest and not the CompletionProvider,
+        // which is currently model based, due to the langauge model.
+        // At some point in the future we should rectify this.
+        let credential = self.credential.read().clone();
+        let request = stream_completion(credential, self.executor.clone(), prompt);
+        async move {
+            let response = request.await?;
+            let stream = response
+                .filter_map(|response| async move {
+                    match response {
+                        Ok(mut response) => Some(Ok(response.choices.pop()?.delta.content?)),
+                        Err(error) => Some(Err(error)),
+                    }
+                })
+                .boxed();
+            Ok(stream)
+        }
+        .boxed()
+    }
+    fn box_clone(&self) -> Box<dyn CompletionProvider> {
+        Box::new((*self).clone())
+    }
+}

crates/ai2/src/providers/open_ai/embedding.rs

@@ -0,0 +1,313 @@
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use futures::AsyncReadExt;
+use gpui2::Executor;
+use gpui2::{serde_json, AppContext};
+use isahc::http::StatusCode;
+use isahc::prelude::Configurable;
+use isahc::{AsyncBody, Response};
+use lazy_static::lazy_static;
+use parking_lot::{Mutex, RwLock};
+use parse_duration::parse;
+use postage::watch;
+use serde::{Deserialize, Serialize};
+use std::env;
+use std::ops::Add;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tiktoken_rs::{cl100k_base, CoreBPE};
+use util::http::{HttpClient, Request};
+use util::ResultExt;
+
+use crate::auth::{CredentialProvider, ProviderCredential};
+use crate::embedding::{Embedding, EmbeddingProvider};
+use crate::models::LanguageModel;
+use crate::providers::open_ai::OpenAILanguageModel;
+
+use crate::providers::open_ai::OPENAI_API_URL;
+
+lazy_static! {
+    static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
+}
+
+#[derive(Clone)]
+pub struct OpenAIEmbeddingProvider {
+    model: OpenAILanguageModel,
+    credential: Arc<RwLock<ProviderCredential>>,
+    pub client: Arc<dyn HttpClient>,
+    pub executor: Arc<Executor>,
+    rate_limit_count_rx: watch::Receiver<Option<Instant>>,
+    rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
+}
+
+#[derive(Serialize)]
+struct OpenAIEmbeddingRequest<'a> {
+    model: &'static str,
+    input: Vec<&'a str>,
+}
+
+#[derive(Deserialize)]
+struct OpenAIEmbeddingResponse {
+    data: Vec<OpenAIEmbedding>,
+    usage: OpenAIEmbeddingUsage,
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenAIEmbedding {
+    embedding: Vec<f32>,
+    index: usize,
+    object: String,
+}
+
+#[derive(Deserialize)]
+struct OpenAIEmbeddingUsage {
+    prompt_tokens: usize,
+    total_tokens: usize,
+}
+
+impl OpenAIEmbeddingProvider {
+    pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Executor>) -> Self {
+        let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
+        let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
+
+        let model = OpenAILanguageModel::load("text-embedding-ada-002");
+        let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
+
+        OpenAIEmbeddingProvider {
+            model,
+            credential,
+            client,
+            executor,
+            rate_limit_count_rx,
+            rate_limit_count_tx,
+        }
+    }
+
+    fn get_api_key(&self) -> Result<String> {
+        match self.credential.read().clone() {
+            ProviderCredential::Credentials { api_key } => Ok(api_key),
+            _ => Err(anyhow!("api credentials not provided")),
+        }
+    }
+
+    fn resolve_rate_limit(&self) {
+        let reset_time = *self.rate_limit_count_tx.lock().borrow();
+
+        if let Some(reset_time) = reset_time {
+            if Instant::now() >= reset_time {
+                *self.rate_limit_count_tx.lock().borrow_mut() = None
+            }
+        }
+
+        log::trace!(
+            "resolving reset time: {:?}",
+            *self.rate_limit_count_tx.lock().borrow()
+        );
+    }
+
+    fn update_reset_time(&self, reset_time: Instant) {
+        let original_time = *self.rate_limit_count_tx.lock().borrow();
+
+        let updated_time = if let Some(original_time) = original_time {
+            if reset_time < original_time {
+                Some(reset_time)
+            } else {
+                Some(original_time)
+            }
+        } else {
+            Some(reset_time)
+        };
+
+        log::trace!("updating rate limit time: {:?}", updated_time);
+
+        *self.rate_limit_count_tx.lock().borrow_mut() = updated_time;
+    }
+    async fn send_request(
+        &self,
+        api_key: &str,
+        spans: Vec<&str>,
+        request_timeout: u64,
+    ) -> Result<Response<AsyncBody>> {
+        let request = Request::post("https://api.openai.com/v1/embeddings")
+            .redirect_policy(isahc::config::RedirectPolicy::Follow)
+            .timeout(Duration::from_secs(request_timeout))
+            .header("Content-Type", "application/json")
+            .header("Authorization", format!("Bearer {}", api_key))
+            .body(
+                serde_json::to_string(&OpenAIEmbeddingRequest {
+                    input: spans.clone(),
+                    model: "text-embedding-ada-002",
+                })
+                .unwrap()
+                .into(),
+            )?;
+
+        Ok(self.client.send(request).await?)
+    }
+}
+
+#[async_trait]
+impl CredentialProvider for OpenAIEmbeddingProvider {
+    fn has_credentials(&self) -> bool {
+        match *self.credential.read() {
+            ProviderCredential::Credentials { .. } => true,
+            _ => false,
+        }
+    }
+    async fn retrieve_credentials(&self, cx: &mut AppContext) -> ProviderCredential {
+        let existing_credential = self.credential.read().clone();
+
+        let retrieved_credential = cx
+            .run_on_main(move |cx| match existing_credential {
+                ProviderCredential::Credentials { .. } => {
+                    return existing_credential.clone();
+                }
+                _ => {
+                    if let Some(api_key) = env::var("OPENAI_API_KEY").log_err() {
+                        return ProviderCredential::Credentials { api_key };
+                    }
+
+                    if let Some(Some((_, api_key))) = cx.read_credentials(OPENAI_API_URL).log_err()
+                    {
+                        if let Some(api_key) = String::from_utf8(api_key).log_err() {
+                            return ProviderCredential::Credentials { api_key };
+                        } else {
+                            return ProviderCredential::NoCredentials;
+                        }
+                    } else {
+                        return ProviderCredential::NoCredentials;
+                    }
+                }
+            })
+            .await;
+
+        *self.credential.write() = retrieved_credential.clone();
+        retrieved_credential
+    }
+
+    async fn save_credentials(&self, cx: &mut AppContext, credential: ProviderCredential) {
+        *self.credential.write() = credential.clone();
+        let credential = credential.clone();
+        cx.run_on_main(move |cx| match credential {
+            ProviderCredential::Credentials { api_key } => {
+                cx.write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
+                    .log_err();
+            }
+            _ => {}
+        })
+        .await;
+    }
+    async fn delete_credentials(&self, cx: &mut AppContext) {
+        cx.run_on_main(move |cx| cx.delete_credentials(OPENAI_API_URL).log_err())
+            .await;
+        *self.credential.write() = ProviderCredential::NoCredentials;
+    }
+}
+
+#[async_trait]
+impl EmbeddingProvider for OpenAIEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(self.model.clone());
+        model
+    }
+
+    fn max_tokens_per_batch(&self) -> usize {
+        50000
+    }
+
+    fn rate_limit_expiration(&self) -> Option<Instant> {
+        *self.rate_limit_count_rx.borrow()
+    }
+
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
+        const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
+        const MAX_RETRIES: usize = 4;
+
+        let api_key = self.get_api_key()?;
+
+        let mut request_number = 0;
+        let mut rate_limiting = false;
+        let mut request_timeout: u64 = 15;
+        let mut response: Response<AsyncBody>;
+        while request_number < MAX_RETRIES {
+            response = self
+                .send_request(
+                    &api_key,
+                    spans.iter().map(|x| &**x).collect(),
+                    request_timeout,
+                )
+                .await?;
+
+            request_number += 1;
+
+            match response.status() {
+                StatusCode::REQUEST_TIMEOUT => {
+                    request_timeout += 5;
+                }
+                StatusCode::OK => {
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+                    let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
+
+                    log::trace!(
+                        "openai embedding completed. tokens: {:?}",
+                        response.usage.total_tokens
+                    );
+
+                    // If we complete a request successfully that was previously rate_limited
+                    // resolve the rate limit
+                    if rate_limiting {
+                        self.resolve_rate_limit()
+                    }
+
+                    return Ok(response
+                        .data
+                        .into_iter()
+                        .map(|embedding| Embedding::from(embedding.embedding))
+                        .collect());
+                }
+                StatusCode::TOO_MANY_REQUESTS => {
+                    rate_limiting = true;
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+
+                    let delay_duration = {
+                        let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
+                        if let Some(time_to_reset) =
+                            response.headers().get("x-ratelimit-reset-tokens")
+                        {
+                            if let Ok(time_str) = time_to_reset.to_str() {
+                                parse(time_str).unwrap_or(delay)
+                            } else {
+                                delay
+                            }
+                        } else {
+                            delay
+                        }
+                    };
+
+                    // If we've previously rate limited, increment the duration but not the count
+                    let reset_time = Instant::now().add(delay_duration);
+                    self.update_reset_time(reset_time);
+
+                    log::trace!(
+                        "openai rate limiting: waiting {:?} until lifted",
+                        &delay_duration
+                    );
+
+                    self.executor.timer(delay_duration).await;
+                }
+                _ => {
+                    let mut body = String::new();
+                    response.body_mut().read_to_string(&mut body).await?;
+                    return Err(anyhow!(
+                        "open ai bad request: {:?} {:?}",
+                        &response.status(),
+                        body
+                    ));
+                }
+            }
+        }
+        Err(anyhow!("openai max retries"))
+    }
+}

crates/ai2/src/providers/open_ai/mod.rs

@@ -0,0 +1,9 @@
+pub mod completion;
+pub mod embedding;
+pub mod model;
+
+pub use completion::*;
+pub use embedding::*;
+pub use model::OpenAILanguageModel;
+
+pub const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";

crates/ai2/src/providers/open_ai/model.rs

@@ -0,0 +1,57 @@
+use anyhow::anyhow;
+use tiktoken_rs::CoreBPE;
+use util::ResultExt;
+
+use crate::models::{LanguageModel, TruncationDirection};
+
+#[derive(Clone)]
+pub struct OpenAILanguageModel {
+    name: String,
+    bpe: Option<CoreBPE>,
+}
+
+impl OpenAILanguageModel {
+    pub fn load(model_name: &str) -> Self {
+        let bpe = tiktoken_rs::get_bpe_from_model(model_name).log_err();
+        OpenAILanguageModel {
+            name: model_name.to_string(),
+            bpe,
+        }
+    }
+}
+
+impl LanguageModel for OpenAILanguageModel {
+    fn name(&self) -> String {
+        self.name.clone()
+    }
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
+        if let Some(bpe) = &self.bpe {
+            anyhow::Ok(bpe.encode_with_special_tokens(content).len())
+        } else {
+            Err(anyhow!("bpe for open ai model was not retrieved"))
+        }
+    }
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String> {
+        if let Some(bpe) = &self.bpe {
+            let tokens = bpe.encode_with_special_tokens(content);
+            if tokens.len() > length {
+                match direction {
+                    TruncationDirection::End => bpe.decode(tokens[..length].to_vec()),
+                    TruncationDirection::Start => bpe.decode(tokens[length..].to_vec()),
+                }
+            } else {
+                bpe.decode(tokens)
+            }
+        } else {
+            Err(anyhow!("bpe for open ai model was not retrieved"))
+        }
+    }
+    fn capacity(&self) -> anyhow::Result<usize> {
+        anyhow::Ok(tiktoken_rs::model::get_context_size(&self.name))
+    }
+}

crates/ai2/src/providers/open_ai/new.rs

@@ -0,0 +1,11 @@
+pub trait LanguageModel {
+    fn name(&self) -> String;
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize>;
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String>;
+    fn capacity(&self) -> anyhow::Result<usize>;
+}

crates/ai2/src/test.rs

@@ -0,0 +1,193 @@
+use std::{
+    sync::atomic::{self, AtomicUsize, Ordering},
+    time::Instant,
+};
+
+use async_trait::async_trait;
+use futures::{channel::mpsc, future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
+use gpui2::AppContext;
+use parking_lot::Mutex;
+
+use crate::{
+    auth::{CredentialProvider, ProviderCredential},
+    completion::{CompletionProvider, CompletionRequest},
+    embedding::{Embedding, EmbeddingProvider},
+    models::{LanguageModel, TruncationDirection},
+};
+
+#[derive(Clone)]
+pub struct FakeLanguageModel {
+    pub capacity: usize,
+}
+
+impl LanguageModel for FakeLanguageModel {
+    fn name(&self) -> String {
+        "dummy".to_string()
+    }
+    fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
+        anyhow::Ok(content.chars().collect::<Vec<char>>().len())
+    }
+    fn truncate(
+        &self,
+        content: &str,
+        length: usize,
+        direction: TruncationDirection,
+    ) -> anyhow::Result<String> {
+        println!("TRYING TO TRUNCATE: {:?}", length.clone());
+
+        if length > self.count_tokens(content)? {
+            println!("NOT TRUNCATING");
+            return anyhow::Ok(content.to_string());
+        }
+
+        anyhow::Ok(match direction {
+            TruncationDirection::End => content.chars().collect::<Vec<char>>()[..length]
+                .into_iter()
+                .collect::<String>(),
+            TruncationDirection::Start => content.chars().collect::<Vec<char>>()[length..]
+                .into_iter()
+                .collect::<String>(),
+        })
+    }
+    fn capacity(&self) -> anyhow::Result<usize> {
+        anyhow::Ok(self.capacity)
+    }
+}
+
+pub struct FakeEmbeddingProvider {
+    pub embedding_count: AtomicUsize,
+}
+
+impl Clone for FakeEmbeddingProvider {
+    fn clone(&self) -> Self {
+        FakeEmbeddingProvider {
+            embedding_count: AtomicUsize::new(self.embedding_count.load(Ordering::SeqCst)),
+        }
+    }
+}
+
+impl Default for FakeEmbeddingProvider {
+    fn default() -> Self {
+        FakeEmbeddingProvider {
+            embedding_count: AtomicUsize::default(),
+        }
+    }
+}
+
+impl FakeEmbeddingProvider {
+    pub fn embedding_count(&self) -> usize {
+        self.embedding_count.load(atomic::Ordering::SeqCst)
+    }
+
+    pub fn embed_sync(&self, span: &str) -> Embedding {
+        let mut result = vec![1.0; 26];
+        for letter in span.chars() {
+            let letter = letter.to_ascii_lowercase();
+            if letter as u32 >= 'a' as u32 {
+                let ix = (letter as u32) - ('a' as u32);
+                if ix < 26 {
+                    result[ix as usize] += 1.0;
+                }
+            }
+        }
+
+        let norm = result.iter().map(|x| x * x).sum::<f32>().sqrt();
+        for x in &mut result {
+            *x /= norm;
+        }
+
+        result.into()
+    }
+}
+
+#[async_trait]
+impl CredentialProvider for FakeEmbeddingProvider {
+    fn has_credentials(&self) -> bool {
+        true
+    }
+    async fn retrieve_credentials(&self, _cx: &mut AppContext) -> ProviderCredential {
+        ProviderCredential::NotNeeded
+    }
+    async fn save_credentials(&self, _cx: &mut AppContext, _credential: ProviderCredential) {}
+    async fn delete_credentials(&self, _cx: &mut AppContext) {}
+}
+
+#[async_trait]
+impl EmbeddingProvider for FakeEmbeddingProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        Box::new(FakeLanguageModel { capacity: 1000 })
+    }
+    fn max_tokens_per_batch(&self) -> usize {
+        1000
+    }
+
+    fn rate_limit_expiration(&self) -> Option<Instant> {
+        None
+    }
+
+    async fn embed_batch(&self, spans: Vec<String>) -> anyhow::Result<Vec<Embedding>> {
+        self.embedding_count
+            .fetch_add(spans.len(), atomic::Ordering::SeqCst);
+
+        anyhow::Ok(spans.iter().map(|span| self.embed_sync(span)).collect())
+    }
+}
+
+pub struct FakeCompletionProvider {
+    last_completion_tx: Mutex<Option<mpsc::Sender<String>>>,
+}
+
+impl Clone for FakeCompletionProvider {
+    fn clone(&self) -> Self {
+        Self {
+            last_completion_tx: Mutex::new(None),
+        }
+    }
+}
+
+impl FakeCompletionProvider {
+    pub fn new() -> Self {
+        Self {
+            last_completion_tx: Mutex::new(None),
+        }
+    }
+
+    pub fn send_completion(&self, completion: impl Into<String>) {
+        let mut tx = self.last_completion_tx.lock();
+        tx.as_mut().unwrap().try_send(completion.into()).unwrap();
+    }
+
+    pub fn finish_completion(&self) {
+        self.last_completion_tx.lock().take().unwrap();
+    }
+}
+
+#[async_trait]
+impl CredentialProvider for FakeCompletionProvider {
+    fn has_credentials(&self) -> bool {
+        true
+    }
+    async fn retrieve_credentials(&self, _cx: &mut AppContext) -> ProviderCredential {
+        ProviderCredential::NotNeeded
+    }
+    async fn save_credentials(&self, _cx: &mut AppContext, _credential: ProviderCredential) {}
+    async fn delete_credentials(&self, _cx: &mut AppContext) {}
+}
+
+impl CompletionProvider for FakeCompletionProvider {
+    fn base_model(&self) -> Box<dyn LanguageModel> {
+        let model: Box<dyn LanguageModel> = Box::new(FakeLanguageModel { capacity: 8190 });
+        model
+    }
+    fn complete(
+        &self,
+        _prompt: Box<dyn CompletionRequest>,
+    ) -> BoxFuture<'static, anyhow::Result<BoxStream<'static, anyhow::Result<String>>>> {
+        let (tx, rx) = mpsc::channel(1);
+        *self.last_completion_tx.lock() = Some(tx);
+        async move { Ok(rx.map(|rx| Ok(rx)).boxed()) }.boxed()
+    }
+    fn box_clone(&self) -> Box<dyn CompletionProvider> {
+        Box::new((*self).clone())
+    }
+}

crates/assistant/Cargo.toml

@@ -17,13 +17,17 @@ fs = { path = "../fs" }
 gpui = { path = "../gpui" }
 language = { path = "../language" }
 menu = { path = "../menu" }
+multi_buffer = { path = "../multi_buffer" }
 search = { path = "../search" }
 settings = { path = "../settings" }
 theme = { path = "../theme" }
 util = { path = "../util" }
 workspace = { path = "../workspace" }
-uuid.workspace = true
+semantic_index = { path = "../semantic_index" }
+project = { path = "../project" }
 
+uuid.workspace = true
+log.workspace = true
 anyhow.workspace = true
 chrono = { version = "0.4", features = ["serde"] }
 futures.workspace = true
@@ -36,11 +40,12 @@ schemars.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 smol.workspace = true
-tiktoken-rs = "0.4"
+tiktoken-rs = "0.5"
 
 [dev-dependencies]
 editor = { path = "../editor", features = ["test-support"] }
 project = { path = "../project", features = ["test-support"] }
+ai = { path = "../ai", features = ["test-support"]}
 
 ctor.workspace = true
 env_logger.workspace = true

crates/assistant/src/assistant.rs

@@ -4,7 +4,7 @@ mod codegen;
 mod prompts;
 mod streaming_diff;
 
-use ai::completion::Role;
+use ai::providers::open_ai::Role;
 use anyhow::Result;
 pub use assistant_panel::AssistantPanel;
 use assistant_settings::OpenAIModel;

crates/assistant/src/codegen.rs

@@ -1,10 +1,11 @@
 use crate::streaming_diff::{Hunk, StreamingDiff};
-use ai::completion::{CompletionProvider, OpenAIRequest};
+use ai::completion::{CompletionProvider, CompletionRequest};
 use anyhow::Result;
-use editor::{multi_buffer, Anchor, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint};
+use editor::{Anchor, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint};
 use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
 use gpui::{Entity, ModelContext, ModelHandle, Task};
 use language::{Rope, TransactionId};
+use multi_buffer;
 use std::{cmp, future, ops::Range, sync::Arc};
 
 pub enum Event {
@@ -95,7 +96,7 @@ impl Codegen {
         self.error.as_ref()
     }
 
-    pub fn start(&mut self, prompt: OpenAIRequest, cx: &mut ModelContext<Self>) {
+    pub fn start(&mut self, prompt: Box<dyn CompletionRequest>, cx: &mut ModelContext<Self>) {
         let range = self.range();
         let snapshot = self.snapshot.clone();
         let selected_text = snapshot
@@ -335,17 +336,25 @@ fn strip_markdown_codeblock(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use futures::{
-        future::BoxFuture,
-        stream::{self, BoxStream},
-    };
+    use ai::test::FakeCompletionProvider;
+    use futures::stream::{self};
     use gpui::{executor::Deterministic, TestAppContext};
     use indoc::indoc;
     use language::{language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, Point};
-    use parking_lot::Mutex;
     use rand::prelude::*;
+    use serde::Serialize;
     use settings::SettingsStore;
-    use smol::future::FutureExt;
+
+    #[derive(Serialize)]
+    pub struct DummyCompletionRequest {
+        pub name: String,
+    }
+
+    impl CompletionRequest for DummyCompletionRequest {
+        fn data(&self) -> serde_json::Result<String> {
+            serde_json::to_string(self)
+        }
+    }
 
     #[gpui::test(iterations = 10)]
     async fn test_transform_autoindent(
@@ -371,7 +380,7 @@ mod tests {
             let snapshot = buffer.snapshot(cx);
             snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
         });
-        let provider = Arc::new(TestCompletionProvider::new());
+        let provider = Arc::new(FakeCompletionProvider::new());
         let codegen = cx.add_model(|cx| {
             Codegen::new(
                 buffer.clone(),
@@ -380,7 +389,11 @@ mod tests {
                 cx,
             )
         });
-        codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
+
+        let request = Box::new(DummyCompletionRequest {
+            name: "test".to_string(),
+        });
+        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
 
         let mut new_text = concat!(
             "       let mut x = 0;\n",
@@ -433,7 +446,7 @@ mod tests {
             let snapshot = buffer.snapshot(cx);
             snapshot.anchor_before(Point::new(1, 6))
         });
-        let provider = Arc::new(TestCompletionProvider::new());
+        let provider = Arc::new(FakeCompletionProvider::new());
         let codegen = cx.add_model(|cx| {
             Codegen::new(
                 buffer.clone(),
@@ -442,7 +455,11 @@ mod tests {
                 cx,
             )
         });
-        codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
+
+        let request = Box::new(DummyCompletionRequest {
+            name: "test".to_string(),
+        });
+        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
 
         let mut new_text = concat!(
             "t mut x = 0;\n",
@@ -495,7 +512,7 @@ mod tests {
             let snapshot = buffer.snapshot(cx);
             snapshot.anchor_before(Point::new(1, 2))
         });
-        let provider = Arc::new(TestCompletionProvider::new());
+        let provider = Arc::new(FakeCompletionProvider::new());
         let codegen = cx.add_model(|cx| {
             Codegen::new(
                 buffer.clone(),
@@ -504,7 +521,11 @@ mod tests {
                 cx,
             )
         });
-        codegen.update(cx, |codegen, cx| codegen.start(Default::default(), cx));
+
+        let request = Box::new(DummyCompletionRequest {
+            name: "test".to_string(),
+        });
+        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
 
         let mut new_text = concat!(
             "let mut x = 0;\n",
@@ -592,38 +613,6 @@ mod tests {
         }
     }
 
-    struct TestCompletionProvider {
-        last_completion_tx: Mutex<Option<mpsc::Sender<String>>>,
-    }
-
-    impl TestCompletionProvider {
-        fn new() -> Self {
-            Self {
-                last_completion_tx: Mutex::new(None),
-            }
-        }
-
-        fn send_completion(&self, completion: impl Into<String>) {
-            let mut tx = self.last_completion_tx.lock();
-            tx.as_mut().unwrap().try_send(completion.into()).unwrap();
-        }
-
-        fn finish_completion(&self) {
-            self.last_completion_tx.lock().take().unwrap();
-        }
-    }
-
-    impl CompletionProvider for TestCompletionProvider {
-        fn complete(
-            &self,
-            _prompt: OpenAIRequest,
-        ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
-            let (tx, rx) = mpsc::channel(1);
-            *self.last_completion_tx.lock() = Some(tx);
-            async move { Ok(rx.map(|rx| Ok(rx)).boxed()) }.boxed()
-        }
-    }
-
     fn rust_lang() -> Language {
         Language::new(
             LanguageConfig {

crates/assistant/src/prompts.rs

@@ -1,8 +1,14 @@
-use crate::codegen::CodegenKind;
+use ai::models::LanguageModel;
+use ai::prompts::base::{PromptArguments, PromptChain, PromptPriority, PromptTemplate};
+use ai::prompts::file_context::FileContext;
+use ai::prompts::generate::GenerateInlineContent;
+use ai::prompts::preamble::EngineerPreamble;
+use ai::prompts::repository_context::{PromptCodeSnippet, RepositoryContext};
+use ai::providers::open_ai::OpenAILanguageModel;
 use language::{BufferSnapshot, OffsetRangeExt, ToOffset};
 use std::cmp::{self, Reverse};
-use std::fmt::Write;
 use std::ops::Range;
+use std::sync::Arc;
 
 #[allow(dead_code)]
 fn summarize(buffer: &BufferSnapshot, selected_range: Range<impl ToOffset>) -> String {
@@ -118,86 +124,50 @@ fn summarize(buffer: &BufferSnapshot, selected_range: Range<impl ToOffset>) -> S
 pub fn generate_content_prompt(
     user_prompt: String,
     language_name: Option<&str>,
-    buffer: &BufferSnapshot,
-    range: Range<impl ToOffset>,
-    kind: CodegenKind,
-) -> String {
-    let range = range.to_offset(buffer);
-    let mut prompt = String::new();
-
-    // General Preamble
-    if let Some(language_name) = language_name {
-        writeln!(prompt, "You're an expert {language_name} engineer.\n").unwrap();
+    buffer: BufferSnapshot,
+    range: Range<usize>,
+    search_results: Vec<PromptCodeSnippet>,
+    model: &str,
+    project_name: Option<String>,
+) -> anyhow::Result<String> {
+    // Using new Prompt Templates
+    let openai_model: Arc<dyn LanguageModel> = Arc::new(OpenAILanguageModel::load(model));
+    let lang_name = if let Some(language_name) = language_name {
+        Some(language_name.to_string())
     } else {
-        writeln!(prompt, "You're an expert engineer.\n").unwrap();
-    }
+        None
+    };
 
-    let mut content = String::new();
-    content.extend(buffer.text_for_range(0..range.start));
-    if range.start == range.end {
-        content.push_str("<|START|>");
-    } else {
-        content.push_str("<|START|");
-    }
-    content.extend(buffer.text_for_range(range.clone()));
-    if range.start != range.end {
-        content.push_str("|END|>");
-    }
-    content.extend(buffer.text_for_range(range.end..buffer.len()));
+    let args = PromptArguments {
+        model: openai_model,
+        language_name: lang_name.clone(),
+        project_name,
+        snippets: search_results.clone(),
+        reserved_tokens: 1000,
+        buffer: Some(buffer),
+        selected_range: Some(range),
+        user_prompt: Some(user_prompt.clone()),
+    };
 
-    writeln!(
-        prompt,
-        "The file you are currently working on has the following content:"
-    )
-    .unwrap();
-    if let Some(language_name) = language_name {
-        let language_name = language_name.to_lowercase();
-        writeln!(prompt, "```{language_name}\n{content}\n```").unwrap();
-    } else {
-        writeln!(prompt, "```\n{content}\n```").unwrap();
-    }
+    let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
+        (PromptPriority::Mandatory, Box::new(EngineerPreamble {})),
+        (
+            PromptPriority::Ordered { order: 1 },
+            Box::new(RepositoryContext {}),
+        ),
+        (
+            PromptPriority::Ordered { order: 0 },
+            Box::new(FileContext {}),
+        ),
+        (
+            PromptPriority::Mandatory,
+            Box::new(GenerateInlineContent {}),
+        ),
+    ];
+    let chain = PromptChain::new(args, templates);
+    let (prompt, _) = chain.generate(true)?;
 
-    match kind {
-        CodegenKind::Generate { position: _ } => {
-            writeln!(prompt, "In particular, the user's cursor is current on the '<|START|>' span in the above outline, with no text selected.").unwrap();
-            writeln!(
-                prompt,
-                "Assume the cursor is located where the `<|START|` marker is."
-            )
-            .unwrap();
-            writeln!(
-                prompt,
-                "Text can't be replaced, so assume your answer will be inserted at the cursor."
-            )
-            .unwrap();
-            writeln!(
-                prompt,
-                "Generate text based on the users prompt: {user_prompt}"
-            )
-            .unwrap();
-        }
-        CodegenKind::Transform { range: _ } => {
-            writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
-            writeln!(
-                prompt,
-                "Modify the users code selected text based upon the users prompt: {user_prompt}"
-            )
-            .unwrap();
-            writeln!(
-                prompt,
-                "You MUST reply with only the adjusted code (within the '<|START|' and '|END|>' spans), not the entire file."
-            )
-            .unwrap();
-        }
-    }
-
-    if let Some(language_name) = language_name {
-        writeln!(prompt, "Your answer MUST always be valid {language_name}").unwrap();
-    }
-    writeln!(prompt, "Always wrap your response in a Markdown codeblock").unwrap();
-    writeln!(prompt, "Never make remarks about the output.").unwrap();
-
-    prompt
+    anyhow::Ok(prompt)
 }
 
 #[cfg(test)]

crates/call/src/call.rs

@@ -10,7 +10,7 @@ use client::{
     ZED_ALWAYS_ACTIVE,
 };
 use collections::HashSet;
-use futures::{future::Shared, FutureExt};
+use futures::{channel::oneshot, future::Shared, Future, FutureExt};
 use gpui::{
     AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Subscription, Task,
     WeakModelHandle,
@@ -37,10 +37,42 @@ pub struct IncomingCall {
     pub initial_project: Option<proto::ParticipantProject>,
 }
 
+pub struct OneAtATime {
+    cancel: Option<oneshot::Sender<()>>,
+}
+
+impl OneAtATime {
+    /// spawn a task in the given context.
+    /// if another task is spawned before that resolves, or if the OneAtATime itself is dropped, the first task will be cancelled and return Ok(None)
+    /// otherwise you'll see the result of the task.
+    fn spawn<F, Fut, R>(&mut self, cx: &mut AppContext, f: F) -> Task<Result<Option<R>>>
+    where
+        F: 'static + FnOnce(AsyncAppContext) -> Fut,
+        Fut: Future<Output = Result<R>>,
+        R: 'static,
+    {
+        let (tx, rx) = oneshot::channel();
+        self.cancel.replace(tx);
+        cx.spawn(|cx| async move {
+            futures::select_biased! {
+                _ = rx.fuse() => Ok(None),
+                result = f(cx).fuse() => result.map(Some),
+            }
+        })
+    }
+
+    fn running(&self) -> bool {
+        self.cancel
+            .as_ref()
+            .is_some_and(|cancel| !cancel.is_canceled())
+    }
+}
+
 /// Singleton global maintaining the user's participation in a room across workspaces.
 pub struct ActiveCall {
     room: Option<(ModelHandle<Room>, Vec<Subscription>)>,
     pending_room_creation: Option<Shared<Task<Result<ModelHandle<Room>, Arc<anyhow::Error>>>>>,
+    _join_debouncer: OneAtATime,
     location: Option<WeakModelHandle<Project>>,
     pending_invites: HashSet<u64>,
     incoming_call: (
@@ -69,6 +101,7 @@ impl ActiveCall {
             pending_invites: Default::default(),
             incoming_call: watch::channel(),
 
+            _join_debouncer: OneAtATime { cancel: None },
             _subscriptions: vec![
                 client.add_request_handler(cx.handle(), Self::handle_incoming_call),
                 client.add_message_handler(cx.handle(), Self::handle_call_canceled),
@@ -143,6 +176,10 @@ impl ActiveCall {
         }
         cx.notify();
 
+        if self._join_debouncer.running() {
+            return Task::ready(Ok(()));
+        }
+
         let room = if let Some(room) = self.room().cloned() {
             Some(Task::ready(Ok(room)).shared())
         } else {
@@ -259,11 +296,20 @@ impl ActiveCall {
             return Task::ready(Err(anyhow!("no incoming call")));
         };
 
-        let join = Room::join(&call, self.client.clone(), self.user_store.clone(), cx);
+        if self.pending_room_creation.is_some() {
+            return Task::ready(Ok(()));
+        }
+
+        let room_id = call.room_id.clone();
+        let client = self.client.clone();
+        let user_store = self.user_store.clone();
+        let join = self
+            ._join_debouncer
+            .spawn(cx, move |cx| Room::join(room_id, client, user_store, cx));
 
         cx.spawn(|this, mut cx| async move {
             let room = join.await?;
-            this.update(&mut cx, |this, cx| this.set_room(Some(room.clone()), cx))
+            this.update(&mut cx, |this, cx| this.set_room(room.clone(), cx))
                 .await?;
             this.update(&mut cx, |this, cx| {
                 this.report_call_event("accept incoming", cx)
@@ -290,20 +336,28 @@ impl ActiveCall {
         &mut self,
         channel_id: u64,
         cx: &mut ModelContext<Self>,
-    ) -> Task<Result<ModelHandle<Room>>> {
+    ) -> Task<Result<Option<ModelHandle<Room>>>> {
         if let Some(room) = self.room().cloned() {
             if room.read(cx).channel_id() == Some(channel_id) {
-                return Task::ready(Ok(room));
+                return Task::ready(Ok(Some(room)));
             } else {
                 room.update(cx, |room, cx| room.clear_state(cx));
             }
         }
 
-        let join = Room::join_channel(channel_id, self.client.clone(), self.user_store.clone(), cx);
+        if self.pending_room_creation.is_some() {
+            return Task::ready(Ok(None));
+        }
 
-        cx.spawn(|this, mut cx| async move {
+        let client = self.client.clone();
+        let user_store = self.user_store.clone();
+        let join = self._join_debouncer.spawn(cx, move |cx| async move {
+            Room::join_channel(channel_id, client, user_store, cx).await
+        });
+
+        cx.spawn(move |this, mut cx| async move {
             let room = join.await?;
-            this.update(&mut cx, |this, cx| this.set_room(Some(room.clone()), cx))
+            this.update(&mut cx, |this, cx| this.set_room(room.clone(), cx))
                 .await?;
             this.update(&mut cx, |this, cx| {
                 this.report_call_event("join channel", cx)
@@ -457,3 +511,40 @@ pub fn report_call_event_for_channel(
     };
     telemetry.report_clickhouse_event(event, telemetry_settings);
 }
+
+#[cfg(test)]
+mod test {
+    use gpui::TestAppContext;
+
+    use crate::OneAtATime;
+
+    #[gpui::test]
+    async fn test_one_at_a_time(cx: &mut TestAppContext) {
+        let mut one_at_a_time = OneAtATime { cancel: None };
+
+        assert_eq!(
+            cx.update(|cx| one_at_a_time.spawn(cx, |_| async { Ok(1) }))
+                .await
+                .unwrap(),
+            Some(1)
+        );
+
+        let (a, b) = cx.update(|cx| {
+            (
+                one_at_a_time.spawn(cx, |_| async {
+                    assert!(false);
+                    Ok(2)
+                }),
+                one_at_a_time.spawn(cx, |_| async { Ok(3) }),
+            )
+        });
+
+        assert_eq!(a.await.unwrap(), None);
+        assert_eq!(b.await.unwrap(), Some(3));
+
+        let promise = cx.update(|cx| one_at_a_time.spawn(cx, |_| async { Ok(4) }));
+        drop(one_at_a_time);
+
+        assert_eq!(promise.await.unwrap(), None);
+    }
+}

crates/call/src/room.rs

@@ -1,7 +1,6 @@
 use crate::{
     call_settings::CallSettings,
     participant::{LocalParticipant, ParticipantLocation, RemoteParticipant, RemoteVideoTrack},
-    IncomingCall,
 };
 use anyhow::{anyhow, Result};
 use audio::{Audio, Sound};
@@ -55,7 +54,7 @@ pub enum Event {
 
 pub struct Room {
     id: u64,
-    channel_id: Option<u64>,
+    pub channel_id: Option<u64>,
     live_kit: Option<LiveKitRoom>,
     status: RoomStatus,
     shared_projects: HashSet<WeakModelHandle<Project>>,
@@ -122,6 +121,10 @@ impl Room {
         }
     }
 
+    pub fn can_publish(&self) -> bool {
+        self.live_kit.as_ref().is_some_and(|room| room.can_publish)
+    }
+
     fn new(
         id: u64,
         channel_id: Option<u64>,
@@ -181,20 +184,23 @@ impl Room {
             });
 
             let connect = room.connect(&connection_info.server_url, &connection_info.token);
-            cx.spawn(|this, mut cx| async move {
-                connect.await?;
+            if connection_info.can_publish {
+                cx.spawn(|this, mut cx| async move {
+                    connect.await?;
 
-                if !cx.read(Self::mute_on_join) {
-                    this.update(&mut cx, |this, cx| this.share_microphone(cx))
-                        .await?;
-                }
+                    if !cx.read(Self::mute_on_join) {
+                        this.update(&mut cx, |this, cx| this.share_microphone(cx))
+                            .await?;
+                    }
 
-                anyhow::Ok(())
-            })
-            .detach_and_log_err(cx);
+                    anyhow::Ok(())
+                })
+                .detach_and_log_err(cx);
+            }
 
             Some(LiveKitRoom {
                 room,
+                can_publish: connection_info.can_publish,
                 screen_track: LocalTrack::None,
                 microphone_track: LocalTrack::None,
                 next_publish_id: 0,
@@ -284,37 +290,32 @@ impl Room {
         })
     }
 
-    pub(crate) fn join_channel(
+    pub(crate) async fn join_channel(
         channel_id: u64,
         client: Arc<Client>,
         user_store: ModelHandle<UserStore>,
-        cx: &mut AppContext,
-    ) -> Task<Result<ModelHandle<Self>>> {
-        cx.spawn(|cx| async move {
-            Self::from_join_response(
-                client.request(proto::JoinChannel { channel_id }).await?,
-                client,
-                user_store,
-                cx,
-            )
-        })
+        cx: AsyncAppContext,
+    ) -> Result<ModelHandle<Self>> {
+        Self::from_join_response(
+            client.request(proto::JoinChannel { channel_id }).await?,
+            client,
+            user_store,
+            cx,
+        )
     }
 
-    pub(crate) fn join(
-        call: &IncomingCall,
+    pub(crate) async fn join(
+        room_id: u64,
         client: Arc<Client>,
         user_store: ModelHandle<UserStore>,
-        cx: &mut AppContext,
-    ) -> Task<Result<ModelHandle<Self>>> {
-        let id = call.room_id;
-        cx.spawn(|cx| async move {
-            Self::from_join_response(
-                client.request(proto::JoinRoom { id }).await?,
-                client,
-                user_store,
-                cx,
-            )
-        })
+        cx: AsyncAppContext,
+    ) -> Result<ModelHandle<Self>> {
+        Self::from_join_response(
+            client.request(proto::JoinRoom { id: room_id }).await?,
+            client,
+            user_store,
+            cx,
+        )
     }
 
     pub fn mute_on_join(cx: &AppContext) -> bool {
@@ -1251,7 +1252,7 @@ impl Room {
                     .read_with(&cx, |this, _| {
                         this.live_kit
                             .as_ref()
-                            .map(|live_kit| live_kit.room.publish_audio_track(&track))
+                            .map(|live_kit| live_kit.room.publish_audio_track(track))
                     })
                     .ok_or_else(|| anyhow!("live-kit was not initialized"))?
                     .await
@@ -1337,7 +1338,7 @@ impl Room {
                     .read_with(&cx, |this, _| {
                         this.live_kit
                             .as_ref()
-                            .map(|live_kit| live_kit.room.publish_video_track(&track))
+                            .map(|live_kit| live_kit.room.publish_video_track(track))
                     })
                     .ok_or_else(|| anyhow!("live-kit was not initialized"))?
                     .await
@@ -1498,6 +1499,7 @@ struct LiveKitRoom {
     deafened: bool,
     speaking: bool,
     next_publish_id: usize,
+    can_publish: bool,
     _maintain_room: Task<()>,
     _maintain_tracks: [Task<()>; 2],
 }

crates/call2/src/call2.rs

@@ -12,8 +12,8 @@ use client2::{
 use collections::HashSet;
 use futures::{future::Shared, FutureExt};
 use gpui2::{
-    AppContext, AsyncAppContext, Context, EventEmitter, Handle, ModelContext, Subscription, Task,
-    WeakHandle,
+    AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, Subscription, Task,
+    WeakModel,
 };
 use postage::watch;
 use project2::Project;
@@ -23,10 +23,10 @@ use std::sync::Arc;
 pub use participant::ParticipantLocation;
 pub use room::Room;
 
-pub fn init(client: Arc<Client>, user_store: Handle<UserStore>, cx: &mut AppContext) {
+pub fn init(client: Arc<Client>, user_store: Model<UserStore>, cx: &mut AppContext) {
     CallSettings::register(cx);
 
-    let active_call = cx.entity(|cx| ActiveCall::new(client, user_store, cx));
+    let active_call = cx.build_model(|cx| ActiveCall::new(client, user_store, cx));
     cx.set_global(active_call);
 }
 
@@ -40,16 +40,16 @@ pub struct IncomingCall {
 
 /// Singleton global maintaining the user's participation in a room across workspaces.
 pub struct ActiveCall {
-    room: Option<(Handle<Room>, Vec<Subscription>)>,
-    pending_room_creation: Option<Shared<Task<Result<Handle<Room>, Arc<anyhow::Error>>>>>,
-    location: Option<WeakHandle<Project>>,
+    room: Option<(Model<Room>, Vec<Subscription>)>,
+    pending_room_creation: Option<Shared<Task<Result<Model<Room>, Arc<anyhow::Error>>>>>,
+    location: Option<WeakModel<Project>>,
     pending_invites: HashSet<u64>,
     incoming_call: (
         watch::Sender<Option<IncomingCall>>,
         watch::Receiver<Option<IncomingCall>>,
     ),
     client: Arc<Client>,
-    user_store: Handle<UserStore>,
+    user_store: Model<UserStore>,
     _subscriptions: Vec<client2::Subscription>,
 }
 
@@ -58,11 +58,7 @@ impl EventEmitter for ActiveCall {
 }
 
 impl ActiveCall {
-    fn new(
-        client: Arc<Client>,
-        user_store: Handle<UserStore>,
-        cx: &mut ModelContext<Self>,
-    ) -> Self {
+    fn new(client: Arc<Client>, user_store: Model<UserStore>, cx: &mut ModelContext<Self>) -> Self {
         Self {
             room: None,
             pending_room_creation: None,
@@ -71,8 +67,8 @@ impl ActiveCall {
             incoming_call: watch::channel(),
 
             _subscriptions: vec![
-                client.add_request_handler(cx.weak_handle(), Self::handle_incoming_call),
-                client.add_message_handler(cx.weak_handle(), Self::handle_call_canceled),
+                client.add_request_handler(cx.weak_model(), Self::handle_incoming_call),
+                client.add_message_handler(cx.weak_model(), Self::handle_call_canceled),
             ],
             client,
             user_store,
@@ -84,7 +80,7 @@ impl ActiveCall {
     }
 
     async fn handle_incoming_call(
-        this: Handle<Self>,
+        this: Model<Self>,
         envelope: TypedEnvelope<proto::IncomingCall>,
         _: Arc<Client>,
         mut cx: AsyncAppContext,
@@ -112,7 +108,7 @@ impl ActiveCall {
     }
 
     async fn handle_call_canceled(
-        this: Handle<Self>,
+        this: Model<Self>,
         envelope: TypedEnvelope<proto::CallCanceled>,
         _: Arc<Client>,
         mut cx: AsyncAppContext,
@@ -129,14 +125,14 @@ impl ActiveCall {
         Ok(())
     }
 
-    pub fn global(cx: &AppContext) -> Handle<Self> {
-        cx.global::<Handle<Self>>().clone()
+    pub fn global(cx: &AppContext) -> Model<Self> {
+        cx.global::<Model<Self>>().clone()
     }
 
     pub fn invite(
         &mut self,
         called_user_id: u64,
-        initial_project: Option<Handle<Project>>,
+        initial_project: Option<Model<Project>>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         if !self.pending_invites.insert(called_user_id) {
@@ -291,7 +287,7 @@ impl ActiveCall {
         &mut self,
         channel_id: u64,
         cx: &mut ModelContext<Self>,
-    ) -> Task<Result<Handle<Room>>> {
+    ) -> Task<Result<Model<Room>>> {
         if let Some(room) = self.room().cloned() {
             if room.read(cx).channel_id() == Some(channel_id) {
                 return Task::ready(Ok(room));
@@ -327,7 +323,7 @@ impl ActiveCall {
 
     pub fn share_project(
         &mut self,
-        project: Handle<Project>,
+        project: Model<Project>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<u64>> {
         if let Some((room, _)) = self.room.as_ref() {
@@ -340,7 +336,7 @@ impl ActiveCall {
 
     pub fn unshare_project(
         &mut self,
-        project: Handle<Project>,
+        project: Model<Project>,
         cx: &mut ModelContext<Self>,
     ) -> Result<()> {
         if let Some((room, _)) = self.room.as_ref() {
@@ -351,13 +347,13 @@ impl ActiveCall {
         }
     }
 
-    pub fn location(&self) -> Option<&WeakHandle<Project>> {
+    pub fn location(&self) -> Option<&WeakModel<Project>> {
         self.location.as_ref()
     }
 
     pub fn set_location(
         &mut self,
-        project: Option<&Handle<Project>>,
+        project: Option<&Model<Project>>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         if project.is_some() || !*ZED_ALWAYS_ACTIVE {
@@ -371,7 +367,7 @@ impl ActiveCall {
 
     fn set_room(
         &mut self,
-        room: Option<Handle<Room>>,
+        room: Option<Model<Room>>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         if room.as_ref() != self.room.as_ref().map(|room| &room.0) {
@@ -407,7 +403,7 @@ impl ActiveCall {
         }
     }
 
-    pub fn room(&self) -> Option<&Handle<Room>> {
+    pub fn room(&self) -> Option<&Model<Room>> {
         self.room.as_ref().map(|(room, _)| room)
     }
 

crates/call2/src/participant.rs

@@ -1,10 +1,8 @@
 use anyhow::{anyhow, Result};
 use client2::ParticipantIndex;
 use client2::{proto, User};
-use collections::HashMap;
-use gpui2::WeakHandle;
+use gpui2::WeakModel;
 pub use live_kit_client::Frame;
-use live_kit_client::RemoteAudioTrack;
 use project2::Project;
 use std::{fmt, sync::Arc};
 
@@ -35,7 +33,7 @@ impl ParticipantLocation {
 #[derive(Clone, Default)]
 pub struct LocalParticipant {
     pub projects: Vec<proto::ParticipantProject>,
-    pub active_project: Option<WeakHandle<Project>>,
+    pub active_project: Option<WeakModel<Project>>,
 }
 
 #[derive(Clone, Debug)]
@@ -47,8 +45,8 @@ pub struct RemoteParticipant {
     pub participant_index: ParticipantIndex,
     pub muted: bool,
     pub speaking: bool,
-    pub video_tracks: HashMap<live_kit_client::Sid, Arc<RemoteVideoTrack>>,
-    pub audio_tracks: HashMap<live_kit_client::Sid, Arc<RemoteAudioTrack>>,
+    // pub video_tracks: HashMap<live_kit_client::Sid, Arc<RemoteVideoTrack>>,
+    // pub audio_tracks: HashMap<live_kit_client::Sid, Arc<RemoteAudioTrack>>,
 }
 
 #[derive(Clone)]

crates/channel/src/channel.rs

@@ -7,10 +7,11 @@ use gpui::{AppContext, ModelHandle};
 use std::sync::Arc;
 
 pub use channel_buffer::{ChannelBuffer, ChannelBufferEvent, ACKNOWLEDGE_DEBOUNCE_INTERVAL};
-pub use channel_chat::{ChannelChat, ChannelChatEvent, ChannelMessage, ChannelMessageId};
-pub use channel_store::{
-    Channel, ChannelData, ChannelEvent, ChannelId, ChannelMembership, ChannelPath, ChannelStore,
+pub use channel_chat::{
+    mentions_to_proto, ChannelChat, ChannelChatEvent, ChannelMessage, ChannelMessageId,
+    MessageParams,
 };
+pub use channel_store::{Channel, ChannelEvent, ChannelId, ChannelMembership, ChannelStore};
 
 #[cfg(test)]
 mod channel_store_tests;

crates/channel/src/channel_buffer.rs

@@ -1,4 +1,4 @@
-use crate::Channel;
+use crate::{Channel, ChannelId, ChannelStore};
 use anyhow::Result;
 use client::{Client, Collaborator, UserStore};
 use collections::HashMap;
@@ -19,10 +19,11 @@ pub(crate) fn init(client: &Arc<Client>) {
 }
 
 pub struct ChannelBuffer {
-    pub(crate) channel: Arc<Channel>,
+    pub channel_id: ChannelId,
     connected: bool,
     collaborators: HashMap<PeerId, Collaborator>,
     user_store: ModelHandle<UserStore>,
+    channel_store: ModelHandle<ChannelStore>,
     buffer: ModelHandle<language::Buffer>,
     buffer_epoch: u64,
     client: Arc<Client>,
@@ -34,6 +35,7 @@ pub enum ChannelBufferEvent {
     CollaboratorsChanged,
     Disconnected,
     BufferEdited,
+    ChannelChanged,
 }
 
 impl Entity for ChannelBuffer {
@@ -46,7 +48,7 @@ impl Entity for ChannelBuffer {
             }
             self.client
                 .send(proto::LeaveChannelBuffer {
-                    channel_id: self.channel.id,
+                    channel_id: self.channel_id,
                 })
                 .log_err();
         }
@@ -58,6 +60,7 @@ impl ChannelBuffer {
         channel: Arc<Channel>,
         client: Arc<Client>,
         user_store: ModelHandle<UserStore>,
+        channel_store: ModelHandle<ChannelStore>,
         mut cx: AsyncAppContext,
     ) -> Result<ModelHandle<Self>> {
         let response = client
@@ -90,9 +93,10 @@ impl ChannelBuffer {
                 connected: true,
                 collaborators: Default::default(),
                 acknowledge_task: None,
-                channel,
+                channel_id: channel.id,
                 subscription: Some(subscription.set_model(&cx.handle(), &mut cx.to_async())),
                 user_store,
+                channel_store,
             };
             this.replace_collaborators(response.collaborators, cx);
             this
@@ -179,7 +183,7 @@ impl ChannelBuffer {
                 let operation = language::proto::serialize_operation(operation);
                 self.client
                     .send(proto::UpdateChannelBuffer {
-                        channel_id: self.channel.id,
+                        channel_id: self.channel_id,
                         operations: vec![operation],
                     })
                     .log_err();
@@ -223,12 +227,15 @@ impl ChannelBuffer {
         &self.collaborators
     }
 
-    pub fn channel(&self) -> Arc<Channel> {
-        self.channel.clone()
+    pub fn channel(&self, cx: &AppContext) -> Option<Arc<Channel>> {
+        self.channel_store
+            .read(cx)
+            .channel_for_id(self.channel_id)
+            .cloned()
     }
 
     pub(crate) fn disconnect(&mut self, cx: &mut ModelContext<Self>) {
-        log::info!("channel buffer {} disconnected", self.channel.id);
+        log::info!("channel buffer {} disconnected", self.channel_id);
         if self.connected {
             self.connected = false;
             self.subscription.take();
@@ -237,6 +244,11 @@ impl ChannelBuffer {
         }
     }
 
+    pub(crate) fn channel_changed(&mut self, cx: &mut ModelContext<Self>) {
+        cx.emit(ChannelBufferEvent::ChannelChanged);
+        cx.notify()
+    }
+
     pub fn is_connected(&self) -> bool {
         self.connected
     }

crates/channel/src/channel_chat.rs

@@ -3,19 +3,25 @@ use anyhow::{anyhow, Result};
 use client::{
     proto,
     user::{User, UserStore},
-    Client, Subscription, TypedEnvelope,
+    Client, Subscription, TypedEnvelope, UserId,
 };
 use futures::lock::Mutex;
 use gpui::{AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task};
 use rand::prelude::*;
-use std::{collections::HashSet, mem, ops::Range, sync::Arc};
+use std::{
+    collections::HashSet,
+    mem,
+    ops::{ControlFlow, Range},
+    sync::Arc,
+};
 use sum_tree::{Bias, SumTree};
 use time::OffsetDateTime;
 use util::{post_inc, ResultExt as _, TryFutureExt};
 
 pub struct ChannelChat {
-    channel: Arc<Channel>,
+    pub channel_id: ChannelId,
     messages: SumTree<ChannelMessage>,
+    acknowledged_message_ids: HashSet<u64>,
     channel_store: ModelHandle<ChannelStore>,
     loaded_all_messages: bool,
     last_acknowledged_id: Option<u64>,
@@ -27,6 +33,12 @@ pub struct ChannelChat {
     _subscription: Subscription,
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub struct MessageParams {
+    pub text: String,
+    pub mentions: Vec<(Range<usize>, UserId)>,
+}
+
 #[derive(Clone, Debug)]
 pub struct ChannelMessage {
     pub id: ChannelMessageId,
@@ -34,6 +46,7 @@ pub struct ChannelMessage {
     pub timestamp: OffsetDateTime,
     pub sender: Arc<User>,
     pub nonce: u128,
+    pub mentions: Vec<(Range<usize>, UserId)>,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -74,7 +87,7 @@ impl Entity for ChannelChat {
     fn release(&mut self, _: &mut AppContext) {
         self.rpc
             .send(proto::LeaveChannelChat {
-                channel_id: self.channel.id,
+                channel_id: self.channel_id,
             })
             .log_err();
     }
@@ -99,12 +112,13 @@ impl ChannelChat {
 
         Ok(cx.add_model(|cx| {
             let mut this = Self {
-                channel,
+                channel_id: channel.id,
                 user_store,
                 channel_store,
                 rpc: client,
                 outgoing_messages_lock: Default::default(),
                 messages: Default::default(),
+                acknowledged_message_ids: Default::default(),
                 loaded_all_messages,
                 next_pending_message_id: 0,
                 last_acknowledged_id: None,
@@ -116,16 +130,23 @@ impl ChannelChat {
         }))
     }
 
-    pub fn channel(&self) -> &Arc<Channel> {
-        &self.channel
+    pub fn channel(&self, cx: &AppContext) -> Option<Arc<Channel>> {
+        self.channel_store
+            .read(cx)
+            .channel_for_id(self.channel_id)
+            .cloned()
+    }
+
+    pub fn client(&self) -> &Arc<Client> {
+        &self.rpc
     }
 
     pub fn send_message(
         &mut self,
-        body: String,
+        message: MessageParams,
         cx: &mut ModelContext<Self>,
-    ) -> Result<Task<Result<()>>> {
-        if body.is_empty() {
+    ) -> Result<Task<Result<u64>>> {
+        if message.text.is_empty() {
             Err(anyhow!("message body can't be empty"))?;
         }
 
@@ -135,16 +156,17 @@ impl ChannelChat {
             .current_user()
             .ok_or_else(|| anyhow!("current_user is not present"))?;
 
-        let channel_id = self.channel.id;
+        let channel_id = self.channel_id;
         let pending_id = ChannelMessageId::Pending(post_inc(&mut self.next_pending_message_id));
         let nonce = self.rng.gen();
         self.insert_messages(
             SumTree::from_item(
                 ChannelMessage {
                     id: pending_id,
-                    body: body.clone(),
+                    body: message.text.clone(),
                     sender: current_user,
                     timestamp: OffsetDateTime::now_utc(),
+                    mentions: message.mentions.clone(),
                     nonce,
                 },
                 &(),
@@ -158,27 +180,25 @@ impl ChannelChat {
             let outgoing_message_guard = outgoing_messages_lock.lock().await;
             let request = rpc.request(proto::SendChannelMessage {
                 channel_id,
-                body,
+                body: message.text,
                 nonce: Some(nonce.into()),
+                mentions: mentions_to_proto(&message.mentions),
             });
             let response = request.await?;
             drop(outgoing_message_guard);
-            let message = ChannelMessage::from_proto(
-                response.message.ok_or_else(|| anyhow!("invalid message"))?,
-                &user_store,
-                &mut cx,
-            )
-            .await?;
+            let response = response.message.ok_or_else(|| anyhow!("invalid message"))?;
+            let id = response.id;
+            let message = ChannelMessage::from_proto(response, &user_store, &mut cx).await?;
             this.update(&mut cx, |this, cx| {
                 this.insert_messages(SumTree::from_item(message, &()), cx);
-                Ok(())
+                Ok(id)
             })
         }))
     }
 
     pub fn remove_message(&mut self, id: u64, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
         let response = self.rpc.request(proto::RemoveChannelMessage {
-            channel_id: self.channel.id,
+            channel_id: self.channel_id,
             message_id: id,
         });
         cx.spawn(|this, mut cx| async move {
@@ -191,41 +211,76 @@ impl ChannelChat {
         })
     }
 
-    pub fn load_more_messages(&mut self, cx: &mut ModelContext<Self>) -> bool {
-        if !self.loaded_all_messages {
-            let rpc = self.rpc.clone();
-            let user_store = self.user_store.clone();
-            let channel_id = self.channel.id;
-            if let Some(before_message_id) =
-                self.messages.first().and_then(|message| match message.id {
-                    ChannelMessageId::Saved(id) => Some(id),
-                    ChannelMessageId::Pending(_) => None,
-                })
-            {
-                cx.spawn(|this, mut cx| {
-                    async move {
-                        let response = rpc
-                            .request(proto::GetChannelMessages {
-                                channel_id,
-                                before_message_id,
-                            })
-                            .await?;
-                        let loaded_all_messages = response.done;
-                        let messages =
-                            messages_from_proto(response.messages, &user_store, &mut cx).await?;
-                        this.update(&mut cx, |this, cx| {
-                            this.loaded_all_messages = loaded_all_messages;
-                            this.insert_messages(messages, cx);
-                        });
-                        anyhow::Ok(())
+    pub fn load_more_messages(&mut self, cx: &mut ModelContext<Self>) -> Option<Task<Option<()>>> {
+        if self.loaded_all_messages {
+            return None;
+        }
+
+        let rpc = self.rpc.clone();
+        let user_store = self.user_store.clone();
+        let channel_id = self.channel_id;
+        let before_message_id = self.first_loaded_message_id()?;
+        Some(cx.spawn(|this, mut cx| {
+            async move {
+                let response = rpc
+                    .request(proto::GetChannelMessages {
+                        channel_id,
+                        before_message_id,
+                    })
+                    .await?;
+                let loaded_all_messages = response.done;
+                let messages = messages_from_proto(response.messages, &user_store, &mut cx).await?;
+                this.update(&mut cx, |this, cx| {
+                    this.loaded_all_messages = loaded_all_messages;
+                    this.insert_messages(messages, cx);
+                });
+                anyhow::Ok(())
+            }
+            .log_err()
+        }))
+    }
+
+    pub fn first_loaded_message_id(&mut self) -> Option<u64> {
+        self.messages.first().and_then(|message| match message.id {
+            ChannelMessageId::Saved(id) => Some(id),
+            ChannelMessageId::Pending(_) => None,
+        })
+    }
+
+    /// Load all of the chat messages since a certain message id.
+    ///
+    /// For now, we always maintain a suffix of the channel's messages.
+    pub async fn load_history_since_message(
+        chat: ModelHandle<Self>,
+        message_id: u64,
+        mut cx: AsyncAppContext,
+    ) -> Option<usize> {
+        loop {
+            let step = chat.update(&mut cx, |chat, cx| {
+                if let Some(first_id) = chat.first_loaded_message_id() {
+                    if first_id <= message_id {
+                        let mut cursor = chat.messages.cursor::<(ChannelMessageId, Count)>();
+                        let message_id = ChannelMessageId::Saved(message_id);
+                        cursor.seek(&message_id, Bias::Left, &());
+                        return ControlFlow::Break(
+                            if cursor
+                                .item()
+                                .map_or(false, |message| message.id == message_id)
+                            {
+                                Some(cursor.start().1 .0)
+                            } else {
+                                None
+                            },
+                        );
                     }
-                    .log_err()
-                })
-                .detach();
-                return true;
+                }
+                ControlFlow::Continue(chat.load_more_messages(cx))
+            });
+            match step {
+                ControlFlow::Break(ix) => return ix,
+                ControlFlow::Continue(task) => task?.await?,
             }
         }
-        false
     }
 
     pub fn acknowledge_last_message(&mut self, cx: &mut ModelContext<Self>) {
@@ -236,13 +291,13 @@ impl ChannelChat {
             {
                 self.rpc
                     .send(proto::AckChannelMessage {
-                        channel_id: self.channel.id,
+                        channel_id: self.channel_id,
                         message_id: latest_message_id,
                     })
                     .ok();
                 self.last_acknowledged_id = Some(latest_message_id);
                 self.channel_store.update(cx, |store, cx| {
-                    store.acknowledge_message_id(self.channel.id, latest_message_id, cx);
+                    store.acknowledge_message_id(self.channel_id, latest_message_id, cx);
                 });
             }
         }
@@ -251,7 +306,7 @@ impl ChannelChat {
     pub fn rejoin(&mut self, cx: &mut ModelContext<Self>) {
         let user_store = self.user_store.clone();
         let rpc = self.rpc.clone();
-        let channel_id = self.channel.id;
+        let channel_id = self.channel_id;
         cx.spawn(|this, mut cx| {
             async move {
                 let response = rpc.request(proto::JoinChannelChat { channel_id }).await?;
@@ -284,6 +339,7 @@ impl ChannelChat {
                     let request = rpc.request(proto::SendChannelMessage {
                         channel_id,
                         body: pending_message.body,
+                        mentions: mentions_to_proto(&pending_message.mentions),
                         nonce: Some(pending_message.nonce.into()),
                     });
                     let response = request.await?;
@@ -319,6 +375,17 @@ impl ChannelChat {
         cursor.item().unwrap()
     }
 
+    pub fn acknowledge_message(&mut self, id: u64) {
+        if self.acknowledged_message_ids.insert(id) {
+            self.rpc
+                .send(proto::AckChannelMessage {
+                    channel_id: self.channel_id,
+                    message_id: id,
+                })
+                .ok();
+        }
+    }
+
     pub fn messages_in_range(&self, range: Range<usize>) -> impl Iterator<Item = &ChannelMessage> {
         let mut cursor = self.messages.cursor::<Count>();
         cursor.seek(&Count(range.start), Bias::Right, &());
@@ -348,7 +415,7 @@ impl ChannelChat {
         this.update(&mut cx, |this, cx| {
             this.insert_messages(SumTree::from_item(message, &()), cx);
             cx.emit(ChannelChatEvent::NewMessage {
-                channel_id: this.channel.id,
+                channel_id: this.channel_id,
                 message_id,
             })
         });
@@ -451,22 +518,7 @@ async fn messages_from_proto(
     user_store: &ModelHandle<UserStore>,
     cx: &mut AsyncAppContext,
 ) -> Result<SumTree<ChannelMessage>> {
-    let unique_user_ids = proto_messages
-        .iter()
-        .map(|m| m.sender_id)
-        .collect::<HashSet<_>>()
-        .into_iter()
-        .collect();
-    user_store
-        .update(cx, |user_store, cx| {
-            user_store.get_users(unique_user_ids, cx)
-        })
-        .await?;
-
-    let mut messages = Vec::with_capacity(proto_messages.len());
-    for message in proto_messages {
-        messages.push(ChannelMessage::from_proto(message, user_store, cx).await?);
-    }
+    let messages = ChannelMessage::from_proto_vec(proto_messages, user_store, cx).await?;
     let mut result = SumTree::new();
     result.extend(messages, &());
     Ok(result)
@@ -486,6 +538,14 @@ impl ChannelMessage {
         Ok(ChannelMessage {
             id: ChannelMessageId::Saved(message.id),
             body: message.body,
+            mentions: message
+                .mentions
+                .into_iter()
+                .filter_map(|mention| {
+                    let range = mention.range?;
+                    Some((range.start as usize..range.end as usize, mention.user_id))
+                })
+                .collect(),
             timestamp: OffsetDateTime::from_unix_timestamp(message.timestamp as i64)?,
             sender,
             nonce: message
@@ -498,6 +558,43 @@ impl ChannelMessage {
     pub fn is_pending(&self) -> bool {
         matches!(self.id, ChannelMessageId::Pending(_))
     }
+
+    pub async fn from_proto_vec(
+        proto_messages: Vec<proto::ChannelMessage>,
+        user_store: &ModelHandle<UserStore>,
+        cx: &mut AsyncAppContext,
+    ) -> Result<Vec<Self>> {
+        let unique_user_ids = proto_messages
+            .iter()
+            .map(|m| m.sender_id)
+            .collect::<HashSet<_>>()
+            .into_iter()
+            .collect();
+        user_store
+            .update(cx, |user_store, cx| {
+                user_store.get_users(unique_user_ids, cx)
+            })
+            .await?;
+
+        let mut messages = Vec::with_capacity(proto_messages.len());
+        for message in proto_messages {
+            messages.push(ChannelMessage::from_proto(message, user_store, cx).await?);
+        }
+        Ok(messages)
+    }
+}
+
+pub fn mentions_to_proto(mentions: &[(Range<usize>, UserId)]) -> Vec<proto::ChatMention> {
+    mentions
+        .iter()
+        .map(|(range, user_id)| proto::ChatMention {
+            range: Some(proto::Range {
+                start: range.start as u64,
+                end: range.end as u64,
+            }),
+            user_id: *user_id as u64,
+        })
+        .collect()
 }
 
 impl sum_tree::Item for ChannelMessage {
@@ -538,3 +635,12 @@ impl<'a> sum_tree::Dimension<'a, ChannelMessageSummary> for Count {
         self.0 += summary.count;
     }
 }
+
+impl<'a> From<&'a str> for MessageParams {
+    fn from(value: &'a str) -> Self {
+        Self {
+            text: value.into(),
+            mentions: Vec::new(),
+        }
+    }
+}

crates/channel/src/channel_store.rs

@@ -1,6 +1,6 @@
 mod channel_index;
 
-use crate::{channel_buffer::ChannelBuffer, channel_chat::ChannelChat};
+use crate::{channel_buffer::ChannelBuffer, channel_chat::ChannelChat, ChannelMessage};
 use anyhow::{anyhow, Result};
 use channel_index::ChannelIndex;
 use client::{Client, Subscription, User, UserId, UserStore};
@@ -9,11 +9,10 @@ use db::RELEASE_CHANNEL;
 use futures::{channel::mpsc, future::Shared, Future, FutureExt, StreamExt};
 use gpui::{AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, WeakModelHandle};
 use rpc::{
-    proto::{self, ChannelEdge, ChannelPermission},
+    proto::{self, ChannelVisibility},
     TypedEnvelope,
 };
-use serde_derive::{Deserialize, Serialize};
-use std::{borrow::Cow, hash::Hash, mem, ops::Deref, sync::Arc, time::Duration};
+use std::{mem, sync::Arc, time::Duration};
 use util::ResultExt;
 
 pub fn init(client: &Arc<Client>, user_store: ModelHandle<UserStore>, cx: &mut AppContext) {
@@ -27,10 +26,9 @@ pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30);
 pub type ChannelId = u64;
 
 pub struct ChannelStore {
-    channel_index: ChannelIndex,
+    pub channel_index: ChannelIndex,
     channel_invitations: Vec<Arc<Channel>>,
     channel_participants: HashMap<ChannelId, Vec<Arc<User>>>,
-    channels_with_admin_privileges: HashSet<ChannelId>,
     outgoing_invites: HashSet<(ChannelId, UserId)>,
     update_channels_tx: mpsc::UnboundedSender<proto::UpdateChannels>,
     opened_buffers: HashMap<ChannelId, OpenedModelHandle<ChannelBuffer>>,
@@ -43,14 +41,15 @@ pub struct ChannelStore {
     _update_channels: Task<()>,
 }
 
-pub type ChannelData = (Channel, ChannelPath);
-
 #[derive(Clone, Debug, PartialEq)]
 pub struct Channel {
     pub id: ChannelId,
     pub name: String,
+    pub visibility: proto::ChannelVisibility,
+    pub role: proto::ChannelRole,
     pub unseen_note_version: Option<(u64, clock::Global)>,
     pub unseen_message_id: Option<u64>,
+    pub parent_path: Vec<u64>,
 }
 
 impl Channel {
@@ -71,15 +70,41 @@ impl Channel {
 
         slug.trim_matches(|c| c == '-').to_string()
     }
-}
 
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
-pub struct ChannelPath(Arc<[ChannelId]>);
+    pub fn can_edit_notes(&self) -> bool {
+        self.role == proto::ChannelRole::Member || self.role == proto::ChannelRole::Admin
+    }
+}
 
 pub struct ChannelMembership {
     pub user: Arc<User>,
     pub kind: proto::channel_member::Kind,
-    pub admin: bool,
+    pub role: proto::ChannelRole,
+}
+impl ChannelMembership {
+    pub fn sort_key(&self) -> MembershipSortKey {
+        MembershipSortKey {
+            role_order: match self.role {
+                proto::ChannelRole::Admin => 0,
+                proto::ChannelRole::Member => 1,
+                proto::ChannelRole::Banned => 2,
+                proto::ChannelRole::Guest => 3,
+            },
+            kind_order: match self.kind {
+                proto::channel_member::Kind::Member => 0,
+                proto::channel_member::Kind::AncestorMember => 1,
+                proto::channel_member::Kind::Invitee => 2,
+            },
+            username_order: self.user.github_login.as_str(),
+        }
+    }
+}
+
+#[derive(PartialOrd, Ord, PartialEq, Eq)]
+pub struct MembershipSortKey<'a> {
+    role_order: u8,
+    kind_order: u8,
+    username_order: &'a str,
 }
 
 pub enum ChannelEvent {
@@ -127,9 +152,6 @@ impl ChannelStore {
                         this.update(&mut cx, |this, cx| this.handle_disconnect(true, cx));
                     }
                 }
-                if status.is_connected() {
-                } else {
-                }
             }
             Some(())
         });
@@ -138,7 +160,6 @@ impl ChannelStore {
             channel_invitations: Vec::default(),
             channel_index: ChannelIndex::default(),
             channel_participants: Default::default(),
-            channels_with_admin_privileges: Default::default(),
             outgoing_invites: Default::default(),
             opened_buffers: Default::default(),
             opened_chats: Default::default(),
@@ -167,16 +188,6 @@ impl ChannelStore {
         self.client.clone()
     }
 
-    pub fn has_children(&self, channel_id: ChannelId) -> bool {
-        self.channel_index.iter().any(|path| {
-            if let Some(ix) = path.iter().position(|id| *id == channel_id) {
-                path.len() > ix + 1
-            } else {
-                false
-            }
-        })
-    }
-
     /// Returns the number of unique channels in the store
     pub fn channel_count(&self) -> usize {
         self.channel_index.by_id().len()
@@ -196,26 +207,31 @@ impl ChannelStore {
     }
 
     /// Iterate over all entries in the channel DAG
-    pub fn channel_dag_entries(&self) -> impl '_ + Iterator<Item = (usize, &Arc<Channel>)> {
-        self.channel_index.iter().map(move |path| {
-            let id = path.last().unwrap();
-            let channel = self.channel_for_id(*id).unwrap();
-            (path.len() - 1, channel)
-        })
+    pub fn ordered_channels(&self) -> impl '_ + Iterator<Item = (usize, &Arc<Channel>)> {
+        self.channel_index
+            .ordered_channels()
+            .iter()
+            .filter_map(move |id| {
+                let channel = self.channel_index.by_id().get(id)?;
+                Some((channel.parent_path.len(), channel))
+            })
     }
 
-    pub fn channel_dag_entry_at(&self, ix: usize) -> Option<(&Arc<Channel>, &ChannelPath)> {
-        let path = self.channel_index.get(ix)?;
-        let id = path.last().unwrap();
-        let channel = self.channel_for_id(*id).unwrap();
-
-        Some((channel, path))
+    pub fn channel_at_index(&self, ix: usize) -> Option<&Arc<Channel>> {
+        let channel_id = self.channel_index.ordered_channels().get(ix)?;
+        self.channel_index.by_id().get(channel_id)
     }
 
     pub fn channel_at(&self, ix: usize) -> Option<&Arc<Channel>> {
         self.channel_index.by_id().values().nth(ix)
     }
 
+    pub fn has_channel_invitation(&self, channel_id: ChannelId) -> bool {
+        self.channel_invitations
+            .iter()
+            .any(|channel| channel.id == channel_id)
+    }
+
     pub fn channel_invitations(&self) -> &[Arc<Channel>] {
         &self.channel_invitations
     }
@@ -240,14 +256,42 @@ impl ChannelStore {
     ) -> Task<Result<ModelHandle<ChannelBuffer>>> {
         let client = self.client.clone();
         let user_store = self.user_store.clone();
+        let channel_store = cx.handle();
         self.open_channel_resource(
             channel_id,
             |this| &mut this.opened_buffers,
-            |channel, cx| ChannelBuffer::new(channel, client, user_store, cx),
+            |channel, cx| ChannelBuffer::new(channel, client, user_store, channel_store, cx),
             cx,
         )
     }
 
+    pub fn fetch_channel_messages(
+        &self,
+        message_ids: Vec<u64>,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<Vec<ChannelMessage>>> {
+        let request = if message_ids.is_empty() {
+            None
+        } else {
+            Some(
+                self.client
+                    .request(proto::GetChannelMessagesById { message_ids }),
+            )
+        };
+        cx.spawn_weak(|this, mut cx| async move {
+            if let Some(request) = request {
+                let response = request.await?;
+                let this = this
+                    .upgrade(&cx)
+                    .ok_or_else(|| anyhow!("channel store dropped"))?;
+                let user_store = this.read_with(&cx, |this, _| this.user_store.clone());
+                ChannelMessage::from_proto_vec(response.messages, &user_store, &mut cx).await
+            } else {
+                Ok(Vec::new())
+            }
+        })
+    }
+
     pub fn has_channel_buffer_changed(&self, channel_id: ChannelId) -> Option<bool> {
         self.channel_index
             .by_id()
@@ -393,16 +437,11 @@ impl ChannelStore {
             .spawn(async move { task.await.map_err(|error| anyhow!("{}", error)) })
     }
 
-    pub fn is_user_admin(&self, channel_id: ChannelId) -> bool {
-        self.channel_index.iter().any(|path| {
-            if let Some(ix) = path.iter().position(|id| *id == channel_id) {
-                path[..=ix]
-                    .iter()
-                    .any(|id| self.channels_with_admin_privileges.contains(id))
-            } else {
-                false
-            }
-        })
+    pub fn is_channel_admin(&self, channel_id: ChannelId) -> bool {
+        let Some(channel) = self.channel_for_id(channel_id) else {
+            return false;
+        };
+        channel.role == proto::ChannelRole::Admin
     }
 
     pub fn channel_participants(&self, channel_id: ChannelId) -> &[Arc<User>] {
@@ -429,24 +468,19 @@ impl ChannelStore {
                 .ok_or_else(|| anyhow!("missing channel in response"))?;
             let channel_id = channel.id;
 
-            let parent_edge = if let Some(parent_id) = parent_id {
-                vec![ChannelEdge {
-                    channel_id: channel.id,
-                    parent_id,
-                }]
-            } else {
-                vec![]
-            };
+            // let parent_edge = if let Some(parent_id) = parent_id {
+            //     vec![ChannelEdge {
+            //         channel_id: channel.id,
+            //         parent_id,
+            //     }]
+            // } else {
+            //     vec![]
+            // };
 
             this.update(&mut cx, |this, cx| {
                 let task = this.update_channels(
                     proto::UpdateChannels {
                         channels: vec![channel],
-                        insert_edge: parent_edge,
-                        channel_permissions: vec![ChannelPermission {
-                            channel_id,
-                            is_admin: true,
-                        }],
                         ..Default::default()
                     },
                     cx,
@@ -464,52 +498,34 @@ impl ChannelStore {
         })
     }
 
-    pub fn link_channel(
-        &mut self,
-        channel_id: ChannelId,
-        to: ChannelId,
-        cx: &mut ModelContext<Self>,
-    ) -> Task<Result<()>> {
-        let client = self.client.clone();
-        cx.spawn(|_, _| async move {
-            let _ = client
-                .request(proto::LinkChannel { channel_id, to })
-                .await?;
-
-            Ok(())
-        })
-    }
-
-    pub fn unlink_channel(
-        &mut self,
-        channel_id: ChannelId,
-        from: ChannelId,
-        cx: &mut ModelContext<Self>,
-    ) -> Task<Result<()>> {
-        let client = self.client.clone();
-        cx.spawn(|_, _| async move {
-            let _ = client
-                .request(proto::UnlinkChannel { channel_id, from })
-                .await?;
-
-            Ok(())
-        })
-    }
-
     pub fn move_channel(
         &mut self,
         channel_id: ChannelId,
-        from: ChannelId,
-        to: ChannelId,
+        to: Option<ChannelId>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         let client = self.client.clone();
         cx.spawn(|_, _| async move {
             let _ = client
-                .request(proto::MoveChannel {
+                .request(proto::MoveChannel { channel_id, to })
+                .await?;
+
+            Ok(())
+        })
+    }
+
+    pub fn set_channel_visibility(
+        &mut self,
+        channel_id: ChannelId,
+        visibility: ChannelVisibility,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<()>> {
+        let client = self.client.clone();
+        cx.spawn(|_, _| async move {
+            let _ = client
+                .request(proto::SetChannelVisibility {
                     channel_id,
-                    from,
-                    to,
+                    visibility: visibility.into(),
                 })
                 .await?;
 
@@ -521,7 +537,7 @@ impl ChannelStore {
         &mut self,
         channel_id: ChannelId,
         user_id: UserId,
-        admin: bool,
+        role: proto::ChannelRole,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         if !self.outgoing_invites.insert((channel_id, user_id)) {
@@ -535,7 +551,7 @@ impl ChannelStore {
                 .request(proto::InviteChannelMember {
                     channel_id,
                     user_id,
-                    admin,
+                    role: role.into(),
                 })
                 .await;
 
@@ -579,11 +595,11 @@ impl ChannelStore {
         })
     }
 
-    pub fn set_member_admin(
+    pub fn set_member_role(
         &mut self,
         channel_id: ChannelId,
         user_id: UserId,
-        admin: bool,
+        role: proto::ChannelRole,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
         if !self.outgoing_invites.insert((channel_id, user_id)) {
@@ -594,10 +610,10 @@ impl ChannelStore {
         let client = self.client.clone();
         cx.spawn(|this, mut cx| async move {
             let result = client
-                .request(proto::SetChannelMemberAdmin {
+                .request(proto::SetChannelMemberRole {
                     channel_id,
                     user_id,
-                    admin,
+                    role: role.into(),
                 })
                 .await;
 
@@ -649,14 +665,15 @@ impl ChannelStore {
         &mut self,
         channel_id: ChannelId,
         accept: bool,
-    ) -> impl Future<Output = Result<()>> {
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<()>> {
         let client = self.client.clone();
-        async move {
+        cx.background().spawn(async move {
             client
                 .request(proto::RespondToChannelInvite { channel_id, accept })
                 .await?;
             Ok(())
-        }
+        })
     }
 
     pub fn get_channel_member_details(
@@ -685,8 +702,8 @@ impl ChannelStore {
                 .filter_map(|(user, member)| {
                     Some(ChannelMembership {
                         user,
-                        admin: member.admin,
-                        kind: proto::channel_member::Kind::from_i32(member.kind)?,
+                        role: member.role(),
+                        kind: member.kind(),
                     })
                 })
                 .collect())
@@ -724,6 +741,11 @@ impl ChannelStore {
     }
 
     fn handle_connect(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
+        self.channel_index.clear();
+        self.channel_invitations.clear();
+        self.channel_participants.clear();
+        self.channel_index.clear();
+        self.outgoing_invites.clear();
         self.disconnect_channel_buffers_task.take();
 
         for chat in self.opened_chats.values() {
@@ -743,7 +765,7 @@ impl ChannelStore {
                     let channel_buffer = buffer.read(cx);
                     let buffer = channel_buffer.buffer().read(cx);
                     buffer_versions.push(proto::ChannelBufferVersion {
-                        channel_id: channel_buffer.channel().id,
+                        channel_id: channel_buffer.channel_id,
                         epoch: channel_buffer.epoch(),
                         version: language::proto::serialize_version(&buffer.version()),
                     });
@@ -770,13 +792,13 @@ impl ChannelStore {
                         };
 
                         channel_buffer.update(cx, |channel_buffer, cx| {
-                            let channel_id = channel_buffer.channel().id;
+                            let channel_id = channel_buffer.channel_id;
                             if let Some(remote_buffer) = response
                                 .buffers
                                 .iter_mut()
                                 .find(|buffer| buffer.channel_id == channel_id)
                             {
-                                let channel_id = channel_buffer.channel().id;
+                                let channel_id = channel_buffer.channel_id;
                                 let remote_version =
                                     language::proto::deserialize_version(&remote_buffer.version);
 
@@ -833,12 +855,6 @@ impl ChannelStore {
     }
 
     fn handle_disconnect(&mut self, wait_for_reconnect: bool, cx: &mut ModelContext<Self>) {
-        self.channel_index.clear();
-        self.channel_invitations.clear();
-        self.channel_participants.clear();
-        self.channels_with_admin_privileges.clear();
-        self.channel_index.clear();
-        self.outgoing_invites.clear();
         cx.notify();
 
         self.disconnect_channel_buffers_task.get_or_insert_with(|| {
@@ -881,9 +897,12 @@ impl ChannelStore {
                     ix,
                     Arc::new(Channel {
                         id: channel.id,
+                        visibility: channel.visibility(),
+                        role: channel.role(),
                         name: channel.name,
                         unseen_note_version: None,
                         unseen_message_id: None,
+                        parent_path: channel.parent_path,
                     }),
                 ),
             }
@@ -891,8 +910,6 @@ impl ChannelStore {
 
         let channels_changed = !payload.channels.is_empty()
             || !payload.delete_channels.is_empty()
-            || !payload.insert_edge.is_empty()
-            || !payload.delete_edge.is_empty()
             || !payload.unseen_channel_messages.is_empty()
             || !payload.unseen_channel_buffer_changes.is_empty();
 
@@ -900,12 +917,17 @@ impl ChannelStore {
             if !payload.delete_channels.is_empty() {
                 self.channel_index.delete_channels(&payload.delete_channels);
                 self.channel_participants
-                    .retain(|channel_id, _| !payload.delete_channels.contains(channel_id));
-                self.channels_with_admin_privileges
-                    .retain(|channel_id| !payload.delete_channels.contains(channel_id));
+                    .retain(|channel_id, _| !&payload.delete_channels.contains(channel_id));
 
                 for channel_id in &payload.delete_channels {
                     let channel_id = *channel_id;
+                    if payload
+                        .channels
+                        .iter()
+                        .any(|channel| channel.id == channel_id)
+                    {
+                        continue;
+                    }
                     if let Some(OpenedModelHandle::Open(buffer)) =
                         self.opened_buffers.remove(&channel_id)
                     {
@@ -918,7 +940,16 @@ impl ChannelStore {
 
             let mut index = self.channel_index.bulk_insert();
             for channel in payload.channels {
-                index.insert(channel)
+                let id = channel.id;
+                let channel_changed = index.insert(channel);
+
+                if channel_changed {
+                    if let Some(OpenedModelHandle::Open(buffer)) = self.opened_buffers.get(&id) {
+                        if let Some(buffer) = buffer.upgrade(cx) {
+                            buffer.update(cx, ChannelBuffer::channel_changed);
+                        }
+                    }
+                }
             }
 
             for unseen_buffer_change in payload.unseen_channel_buffer_changes {
@@ -936,24 +967,6 @@ impl ChannelStore {
                     unseen_channel_message.message_id,
                 );
             }
-
-            for edge in payload.insert_edge {
-                index.insert_edge(edge.channel_id, edge.parent_id);
-            }
-
-            for edge in payload.delete_edge {
-                index.delete_edge(edge.parent_id, edge.channel_id);
-            }
-        }
-
-        for permission in payload.channel_permissions {
-            if permission.is_admin {
-                self.channels_with_admin_privileges
-                    .insert(permission.channel_id);
-            } else {
-                self.channels_with_admin_privileges
-                    .remove(&permission.channel_id);
-            }
         }
 
         cx.notify();
@@ -1002,44 +1015,3 @@ impl ChannelStore {
         }))
     }
 }
-
-impl Deref for ChannelPath {
-    type Target = [ChannelId];
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl ChannelPath {
-    pub fn new(path: Arc<[ChannelId]>) -> Self {
-        debug_assert!(path.len() >= 1);
-        Self(path)
-    }
-
-    pub fn parent_id(&self) -> Option<ChannelId> {
-        self.0.len().checked_sub(2).map(|i| self.0[i])
-    }
-
-    pub fn channel_id(&self) -> ChannelId {
-        self.0[self.0.len() - 1]
-    }
-}
-
-impl From<ChannelPath> for Cow<'static, ChannelPath> {
-    fn from(value: ChannelPath) -> Self {
-        Cow::Owned(value)
-    }
-}
-
-impl<'a> From<&'a ChannelPath> for Cow<'a, ChannelPath> {
-    fn from(value: &'a ChannelPath) -> Self {
-        Cow::Borrowed(value)
-    }
-}
-
-impl Default for ChannelPath {
-    fn default() -> Self {
-        ChannelPath(Arc::from([]))
-    }
-}

crates/channel/src/channel_store/channel_index.rs

@@ -1,14 +1,11 @@
-use std::{ops::Deref, sync::Arc};
-
 use crate::{Channel, ChannelId};
 use collections::BTreeMap;
 use rpc::proto;
-
-use super::ChannelPath;
+use std::sync::Arc;
 
 #[derive(Default, Debug)]
 pub struct ChannelIndex {
-    paths: Vec<ChannelPath>,
+    channels_ordered: Vec<ChannelId>,
     channels_by_id: BTreeMap<ChannelId, Arc<Channel>>,
 }
 
@@ -17,8 +14,12 @@ impl ChannelIndex {
         &self.channels_by_id
     }
 
+    pub fn ordered_channels(&self) -> &[ChannelId] {
+        &self.channels_ordered
+    }
+
     pub fn clear(&mut self) {
-        self.paths.clear();
+        self.channels_ordered.clear();
         self.channels_by_id.clear();
     }
 
@@ -26,15 +27,13 @@ impl ChannelIndex {
     pub fn delete_channels(&mut self, channels: &[ChannelId]) {
         self.channels_by_id
             .retain(|channel_id, _| !channels.contains(channel_id));
-        self.paths.retain(|path| {
-            path.iter()
-                .all(|channel_id| self.channels_by_id.contains_key(channel_id))
-        });
+        self.channels_ordered
+            .retain(|channel_id| !channels.contains(channel_id));
     }
 
     pub fn bulk_insert(&mut self) -> ChannelPathsInsertGuard {
         ChannelPathsInsertGuard {
-            paths: &mut self.paths,
+            channels_ordered: &mut self.channels_ordered,
             channels_by_id: &mut self.channels_by_id,
         }
     }
@@ -77,42 +76,15 @@ impl ChannelIndex {
     }
 }
 
-impl Deref for ChannelIndex {
-    type Target = [ChannelPath];
-
-    fn deref(&self) -> &Self::Target {
-        &self.paths
-    }
-}
-
 /// A guard for ensuring that the paths index maintains its sort and uniqueness
 /// invariants after a series of insertions
 #[derive(Debug)]
 pub struct ChannelPathsInsertGuard<'a> {
-    paths: &'a mut Vec<ChannelPath>,
+    channels_ordered: &'a mut Vec<ChannelId>,
     channels_by_id: &'a mut BTreeMap<ChannelId, Arc<Channel>>,
 }
 
 impl<'a> ChannelPathsInsertGuard<'a> {
-    /// Remove the given edge from this index. This will not remove the channel.
-    /// If this operation would result in a dangling edge, re-insert it.
-    pub fn delete_edge(&mut self, parent_id: ChannelId, channel_id: ChannelId) {
-        self.paths.retain(|path| {
-            !path
-                .windows(2)
-                .any(|window| window == [parent_id, channel_id])
-        });
-
-        // Ensure that there is at least one channel path in the index
-        if !self
-            .paths
-            .iter()
-            .any(|path| path.iter().any(|id| id == &channel_id))
-        {
-            self.insert_root(channel_id);
-        }
-    }
-
     pub fn note_changed(&mut self, channel_id: ChannelId, epoch: u64, version: &clock::Global) {
         insert_note_changed(&mut self.channels_by_id, channel_id, epoch, &version);
     }
@@ -121,91 +93,65 @@ impl<'a> ChannelPathsInsertGuard<'a> {
         insert_new_message(&mut self.channels_by_id, channel_id, message_id)
     }
 
-    pub fn insert(&mut self, channel_proto: proto::Channel) {
+    pub fn insert(&mut self, channel_proto: proto::Channel) -> bool {
+        let mut ret = false;
         if let Some(existing_channel) = self.channels_by_id.get_mut(&channel_proto.id) {
-            Arc::make_mut(existing_channel).name = channel_proto.name;
+            let existing_channel = Arc::make_mut(existing_channel);
+
+            ret = existing_channel.visibility != channel_proto.visibility()
+                || existing_channel.role != channel_proto.role()
+                || existing_channel.name != channel_proto.name;
+
+            existing_channel.visibility = channel_proto.visibility();
+            existing_channel.role = channel_proto.role();
+            existing_channel.name = channel_proto.name;
         } else {
             self.channels_by_id.insert(
                 channel_proto.id,
                 Arc::new(Channel {
                     id: channel_proto.id,
+                    visibility: channel_proto.visibility(),
+                    role: channel_proto.role(),
                     name: channel_proto.name,
                     unseen_note_version: None,
                     unseen_message_id: None,
+                    parent_path: channel_proto.parent_path,
                 }),
             );
             self.insert_root(channel_proto.id);
         }
-    }
-
-    pub fn insert_edge(&mut self, channel_id: ChannelId, parent_id: ChannelId) {
-        let mut parents = Vec::new();
-        let mut descendants = Vec::new();
-        let mut ixs_to_remove = Vec::new();
-
-        for (ix, path) in self.paths.iter().enumerate() {
-            if path
-                .windows(2)
-                .any(|window| window[0] == parent_id && window[1] == channel_id)
-            {
-                // We already have this edge in the index
-                return;
-            }
-            if path.ends_with(&[parent_id]) {
-                parents.push(path);
-            } else if let Some(position) = path.iter().position(|id| id == &channel_id) {
-                if position == 0 {
-                    ixs_to_remove.push(ix);
-                }
-                descendants.push(path.split_at(position).1);
-            }
-        }
-
-        let mut new_paths = Vec::new();
-        for parent in parents.iter() {
-            if descendants.is_empty() {
-                let mut new_path = Vec::with_capacity(parent.len() + 1);
-                new_path.extend_from_slice(parent);
-                new_path.push(channel_id);
-                new_paths.push(ChannelPath::new(new_path.into()));
-            } else {
-                for descendant in descendants.iter() {
-                    let mut new_path = Vec::with_capacity(parent.len() + descendant.len());
-                    new_path.extend_from_slice(parent);
-                    new_path.extend_from_slice(descendant);
-                    new_paths.push(ChannelPath::new(new_path.into()));
-                }
-            }
-        }
-
-        for ix in ixs_to_remove.into_iter().rev() {
-            self.paths.swap_remove(ix);
-        }
-        self.paths.extend(new_paths)
+        ret
     }
 
     fn insert_root(&mut self, channel_id: ChannelId) {
-        self.paths.push(ChannelPath::new(Arc::from([channel_id])));
+        self.channels_ordered.push(channel_id);
     }
 }
 
 impl<'a> Drop for ChannelPathsInsertGuard<'a> {
     fn drop(&mut self) {
-        self.paths.sort_by(|a, b| {
-            let a = channel_path_sorting_key(a, &self.channels_by_id);
-            let b = channel_path_sorting_key(b, &self.channels_by_id);
+        self.channels_ordered.sort_by(|a, b| {
+            let a = channel_path_sorting_key(*a, &self.channels_by_id);
+            let b = channel_path_sorting_key(*b, &self.channels_by_id);
             a.cmp(b)
         });
-        self.paths.dedup();
+        self.channels_ordered.dedup();
     }
 }
 
 fn channel_path_sorting_key<'a>(
-    path: &'a [ChannelId],
+    id: ChannelId,
     channels_by_id: &'a BTreeMap<ChannelId, Arc<Channel>>,
-) -> impl 'a + Iterator<Item = Option<&'a str>> {
-    path.iter()
-        .map(|id| Some(channels_by_id.get(id)?.name.as_str()))
+) -> impl Iterator<Item = &str> {
+    let (parent_path, name) = channels_by_id
+        .get(&id)
+        .map_or((&[] as &[_], None), |channel| {
+            (channel.parent_path.as_slice(), Some(channel.name.as_str()))
+        });
+    parent_path
+        .iter()
+        .filter_map(|id| Some(channels_by_id.get(id)?.name.as_str()))
+        .chain(name)
 }
 
 fn insert_note_changed(

crates/channel/src/channel_store_tests.rs

@@ -3,7 +3,7 @@ use crate::channel_chat::ChannelChatEvent;
 use super::*;
 use client::{test::FakeServer, Client, UserStore};
 use gpui::{AppContext, ModelHandle, TestAppContext};
-use rpc::proto;
+use rpc::proto::{self};
 use settings::SettingsStore;
 use util::http::FakeHttpClient;
 
@@ -18,16 +18,18 @@ fn test_update_channels(cx: &mut AppContext) {
                 proto::Channel {
                     id: 1,
                     name: "b".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Admin.into(),
+                    parent_path: Vec::new(),
                 },
                 proto::Channel {
                     id: 2,
                     name: "a".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Member.into(),
+                    parent_path: Vec::new(),
                 },
             ],
-            channel_permissions: vec![proto::ChannelPermission {
-                channel_id: 1,
-                is_admin: true,
-            }],
             ..Default::default()
         },
         cx,
@@ -36,8 +38,8 @@ fn test_update_channels(cx: &mut AppContext) {
         &channel_store,
         &[
             //
-            (0, "a".to_string(), false),
-            (0, "b".to_string(), true),
+            (0, "a".to_string(), proto::ChannelRole::Member),
+            (0, "b".to_string(), proto::ChannelRole::Admin),
         ],
         cx,
     );
@@ -49,20 +51,16 @@ fn test_update_channels(cx: &mut AppContext) {
                 proto::Channel {
                     id: 3,
                     name: "x".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Admin.into(),
+                    parent_path: vec![1],
                 },
                 proto::Channel {
                     id: 4,
                     name: "y".to_string(),
-                },
-            ],
-            insert_edge: vec![
-                proto::ChannelEdge {
-                    parent_id: 1,
-                    channel_id: 3,
-                },
-                proto::ChannelEdge {
-                    parent_id: 2,
-                    channel_id: 4,
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Member.into(),
+                    parent_path: vec![2],
                 },
             ],
             ..Default::default()
@@ -72,10 +70,10 @@ fn test_update_channels(cx: &mut AppContext) {
     assert_channels(
         &channel_store,
         &[
-            (0, "a".to_string(), false),
-            (1, "y".to_string(), false),
-            (0, "b".to_string(), true),
-            (1, "x".to_string(), true),
+            (0, "a".to_string(), proto::ChannelRole::Member),
+            (1, "y".to_string(), proto::ChannelRole::Member),
+            (0, "b".to_string(), proto::ChannelRole::Admin),
+            (1, "x".to_string(), proto::ChannelRole::Admin),
         ],
         cx,
     );
@@ -92,30 +90,25 @@ fn test_dangling_channel_paths(cx: &mut AppContext) {
                 proto::Channel {
                     id: 0,
                     name: "a".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Admin.into(),
+                    parent_path: vec![],
                 },
                 proto::Channel {
                     id: 1,
                     name: "b".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Admin.into(),
+                    parent_path: vec![0],
                 },
                 proto::Channel {
                     id: 2,
                     name: "c".to_string(),
+                    visibility: proto::ChannelVisibility::Members as i32,
+                    role: proto::ChannelRole::Admin.into(),
+                    parent_path: vec![0, 1],
                 },
             ],
-            insert_edge: vec![
-                proto::ChannelEdge {
-                    parent_id: 0,
-                    channel_id: 1,
-                },
-                proto::ChannelEdge {
-                    parent_id: 1,
-                    channel_id: 2,
-                },
-            ],
-            channel_permissions: vec![proto::ChannelPermission {
-                channel_id: 0,
-                is_admin: true,
-            }],
             ..Default::default()
         },
         cx,
@@ -125,9 +118,9 @@ fn test_dangling_channel_paths(cx: &mut AppContext) {
         &channel_store,
         &[
             //
-            (0, "a".to_string(), true),
-            (1, "b".to_string(), true),
-            (2, "c".to_string(), true),
+            (0, "a".to_string(), proto::ChannelRole::Admin),
+            (1, "b".to_string(), proto::ChannelRole::Admin),
+            (2, "c".to_string(), proto::ChannelRole::Admin),
         ],
         cx,
     );
@@ -142,7 +135,11 @@ fn test_dangling_channel_paths(cx: &mut AppContext) {
     );
 
     // Make sure that the 1/2/3 path is gone
-    assert_channels(&channel_store, &[(0, "a".to_string(), true)], cx);
+    assert_channels(
+        &channel_store,
+        &[(0, "a".to_string(), proto::ChannelRole::Admin)],
+        cx,
+    );
 }
 
 #[gpui::test]
@@ -158,12 +155,19 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
         channels: vec![proto::Channel {
             id: channel_id,
             name: "the-channel".to_string(),
+            visibility: proto::ChannelVisibility::Members as i32,
+            role: proto::ChannelRole::Member.into(),
+            parent_path: vec![],
         }],
         ..Default::default()
     });
     cx.foreground().run_until_parked();
     cx.read(|cx| {
-        assert_channels(&channel_store, &[(0, "the-channel".to_string(), false)], cx);
+        assert_channels(
+            &channel_store,
+            &[(0, "the-channel".to_string(), proto::ChannelRole::Member)],
+            cx,
+        );
     });
 
     let get_users = server.receive::<proto::GetUsers>().await.unwrap();
@@ -181,7 +185,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
 
     // Join a channel and populate its existing messages.
     let channel = channel_store.update(cx, |store, cx| {
-        let channel_id = store.channel_dag_entries().next().unwrap().1.id;
+        let channel_id = store.ordered_channels().next().unwrap().1.id;
         store.open_channel_chat(channel_id, cx)
     });
     let join_channel = server.receive::<proto::JoinChannelChat>().await.unwrap();
@@ -194,6 +198,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                     body: "a".into(),
                     timestamp: 1000,
                     sender_id: 5,
+                    mentions: vec![],
                     nonce: Some(1.into()),
                 },
                 proto::ChannelMessage {
@@ -201,6 +206,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                     body: "b".into(),
                     timestamp: 1001,
                     sender_id: 6,
+                    mentions: vec![],
                     nonce: Some(2.into()),
                 },
             ],
@@ -247,6 +253,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
             body: "c".into(),
             timestamp: 1002,
             sender_id: 7,
+            mentions: vec![],
             nonce: Some(3.into()),
         }),
     });
@@ -284,7 +291,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
 
     // Scroll up to view older messages.
     channel.update(cx, |channel, cx| {
-        assert!(channel.load_more_messages(cx));
+        channel.load_more_messages(cx).unwrap().detach();
     });
     let get_messages = server.receive::<proto::GetChannelMessages>().await.unwrap();
     assert_eq!(get_messages.payload.channel_id, 5);
@@ -300,6 +307,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                     timestamp: 998,
                     sender_id: 5,
                     nonce: Some(4.into()),
+                    mentions: vec![],
                 },
                 proto::ChannelMessage {
                     id: 9,
@@ -307,6 +315,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                     timestamp: 999,
                     sender_id: 6,
                     nonce: Some(5.into()),
+                    mentions: vec![],
                 },
             ],
         },
@@ -358,19 +367,13 @@ fn update_channels(
 #[track_caller]
 fn assert_channels(
     channel_store: &ModelHandle<ChannelStore>,
-    expected_channels: &[(usize, String, bool)],
+    expected_channels: &[(usize, String, proto::ChannelRole)],
     cx: &AppContext,
 ) {
     let actual = channel_store.read_with(cx, |store, _| {
         store
-            .channel_dag_entries()
-            .map(|(depth, channel)| {
-                (
-                    depth,
-                    channel.name.to_string(),
-                    store.is_user_admin(channel.id),
-                )
-            })
+            .ordered_channels()
+            .map(|(depth, channel)| (depth, channel.name.to_string(), channel.role))
             .collect::<Vec<_>>()
     });
     assert_eq!(actual, expected_channels);

crates/client/src/telemetry.rs

@@ -4,7 +4,9 @@ use lazy_static::lazy_static;
 use parking_lot::Mutex;
 use serde::Serialize;
 use std::{env, io::Write, mem, path::PathBuf, sync::Arc, time::Duration};
-use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt};
+use sysinfo::{
+    CpuRefreshKind, Pid, PidExt, ProcessExt, ProcessRefreshKind, RefreshKind, System, SystemExt,
+};
 use tempfile::NamedTempFile;
 use util::http::HttpClient;
 use util::{channel::ReleaseChannel, TryFutureExt};
@@ -166,8 +168,16 @@ impl Telemetry {
 
         let this = self.clone();
         cx.spawn(|mut cx| async move {
-            let mut system = System::new_all();
-            system.refresh_all();
+            // Avoiding calling `System::new_all()`, as there have been crashes related to it
+            let refresh_kind = RefreshKind::new()
+                .with_memory() // For memory usage
+                .with_processes(ProcessRefreshKind::everything()) // For process usage
+                .with_cpu(CpuRefreshKind::everything()); // For core count
+
+            let mut system = System::new_with_specifics(refresh_kind);
+
+            // Avoiding calling `refresh_all()`, just update what we need
+            system.refresh_specifics(refresh_kind);
 
             loop {
                 // Waiting some amount of time before the first query is important to get a reasonable value
@@ -175,8 +185,7 @@ impl Telemetry {
                 const DURATION_BETWEEN_SYSTEM_EVENTS: Duration = Duration::from_secs(60);
                 smol::Timer::after(DURATION_BETWEEN_SYSTEM_EVENTS).await;
 
-                system.refresh_memory();
-                system.refresh_processes();
+                system.refresh_specifics(refresh_kind);
 
                 let current_process = Pid::from_u32(std::process::id());
                 let Some(process) = system.processes().get(&current_process) else {

crates/client/src/user.rs

@@ -293,21 +293,19 @@ impl UserStore {
                     // No need to paralellize here
                     let mut updated_contacts = Vec::new();
                     for contact in message.contacts {
-                        let should_notify = contact.should_notify;
-                        updated_contacts.push((
-                            Arc::new(Contact::from_proto(contact, &this, &mut cx).await?),
-                            should_notify,
+                        updated_contacts.push(Arc::new(
+                            Contact::from_proto(contact, &this, &mut cx).await?,
                         ));
                     }
 
                     let mut incoming_requests = Vec::new();
                     for request in message.incoming_requests {
-                        incoming_requests.push({
-                            let user = this
-                                .update(&mut cx, |this, cx| this.get_user(request.requester_id, cx))
-                                .await?;
-                            (user, request.should_notify)
-                        });
+                        incoming_requests.push(
+                            this.update(&mut cx, |this, cx| {
+                                this.get_user(request.requester_id, cx)
+                            })
+                            .await?,
+                        );
                     }
 
                     let mut outgoing_requests = Vec::new();
@@ -330,13 +328,7 @@ impl UserStore {
                         this.contacts
                             .retain(|contact| !removed_contacts.contains(&contact.user.id));
                         // Update existing contacts and insert new ones
-                        for (updated_contact, should_notify) in updated_contacts {
-                            if should_notify {
-                                cx.emit(Event::Contact {
-                                    user: updated_contact.user.clone(),
-                                    kind: ContactEventKind::Accepted,
-                                });
-                            }
+                        for updated_contact in updated_contacts {
                             match this.contacts.binary_search_by_key(
                                 &&updated_contact.user.github_login,
                                 |contact| &contact.user.github_login,
@@ -359,14 +351,7 @@ impl UserStore {
                             }
                         });
                         // Update existing incoming requests and insert new ones
-                        for (user, should_notify) in incoming_requests {
-                            if should_notify {
-                                cx.emit(Event::Contact {
-                                    user: user.clone(),
-                                    kind: ContactEventKind::Requested,
-                                });
-                            }
-
+                        for user in incoming_requests {
                             match this
                                 .incoming_contact_requests
                                 .binary_search_by_key(&&user.github_login, |contact| {
@@ -415,6 +400,12 @@ impl UserStore {
         &self.incoming_contact_requests
     }
 
+    pub fn has_incoming_contact_request(&self, user_id: u64) -> bool {
+        self.incoming_contact_requests
+            .iter()
+            .any(|user| user.id == user_id)
+    }
+
     pub fn outgoing_contact_requests(&self) -> &[Arc<User>] {
         &self.outgoing_contact_requests
     }

crates/client2/src/client2.rs

@@ -14,8 +14,8 @@ use futures::{
     future::BoxFuture, AsyncReadExt, FutureExt, SinkExt, StreamExt, TryFutureExt as _, TryStreamExt,
 };
 use gpui2::{
-    serde_json, AnyHandle, AnyWeakHandle, AppContext, AsyncAppContext, Handle, SemanticVersion,
-    Task, WeakHandle,
+    serde_json, AnyModel, AnyWeakModel, AppContext, AsyncAppContext, Model, SemanticVersion, Task,
+    WeakModel,
 };
 use lazy_static::lazy_static;
 use parking_lot::RwLock;
@@ -227,7 +227,7 @@ struct ClientState {
     _reconnect_task: Option<Task<()>>,
     reconnect_interval: Duration,
     entities_by_type_and_remote_id: HashMap<(TypeId, u64), WeakSubscriber>,
-    models_by_message_type: HashMap<TypeId, AnyWeakHandle>,
+    models_by_message_type: HashMap<TypeId, AnyWeakModel>,
     entity_types_by_message_type: HashMap<TypeId, TypeId>,
     #[allow(clippy::type_complexity)]
     message_handlers: HashMap<
@@ -236,7 +236,7 @@ struct ClientState {
             dyn Send
                 + Sync
                 + Fn(
-                    AnyHandle,
+                    AnyModel,
                     Box<dyn AnyTypedEnvelope>,
                     &Arc<Client>,
                     AsyncAppContext,
@@ -246,7 +246,7 @@ struct ClientState {
 }
 
 enum WeakSubscriber {
-    Entity { handle: AnyWeakHandle },
+    Entity { handle: AnyWeakModel },
     Pending(Vec<Box<dyn AnyTypedEnvelope>>),
 }
 
@@ -312,9 +312,9 @@ pub struct PendingEntitySubscription<T: 'static> {
 
 impl<T> PendingEntitySubscription<T>
 where
-    T: 'static + Send + Sync,
+    T: 'static + Send,
 {
-    pub fn set_model(mut self, model: &Handle<T>, cx: &mut AsyncAppContext) -> Subscription {
+    pub fn set_model(mut self, model: &Model<T>, cx: &mut AsyncAppContext) -> Subscription {
         self.consumed = true;
         let mut state = self.client.state.write();
         let id = (TypeId::of::<T>(), self.remote_id);
@@ -529,7 +529,7 @@ impl Client {
         remote_id: u64,
     ) -> Result<PendingEntitySubscription<T>>
     where
-        T: 'static + Send + Sync,
+        T: 'static + Send,
     {
         let id = (TypeId::of::<T>(), remote_id);
 
@@ -552,13 +552,13 @@ impl Client {
     #[track_caller]
     pub fn add_message_handler<M, E, H, F>(
         self: &Arc<Self>,
-        entity: WeakHandle<E>,
+        entity: WeakModel<E>,
         handler: H,
     ) -> Subscription
     where
         M: EnvelopedMessage,
-        E: 'static + Send + Sync,
-        H: 'static + Send + Sync + Fn(Handle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
+        E: 'static + Send,
+        H: 'static + Send + Sync + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
         F: 'static + Future<Output = Result<()>> + Send,
     {
         let message_type_id = TypeId::of::<M>();
@@ -594,13 +594,13 @@ impl Client {
 
     pub fn add_request_handler<M, E, H, F>(
         self: &Arc<Self>,
-        model: WeakHandle<E>,
+        model: WeakModel<E>,
         handler: H,
     ) -> Subscription
     where
         M: RequestMessage,
-        E: 'static + Send + Sync,
-        H: 'static + Send + Sync + Fn(Handle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
+        E: 'static + Send,
+        H: 'static + Send + Sync + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
         F: 'static + Future<Output = Result<M::Response>> + Send,
     {
         self.add_message_handler(model, move |handle, envelope, this, cx| {
@@ -615,8 +615,8 @@ impl Client {
     pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
     where
         M: EntityMessage,
-        E: 'static + Send + Sync,
-        H: 'static + Send + Sync + Fn(Handle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
+        E: 'static + Send,
+        H: 'static + Send + Sync + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
         F: 'static + Future<Output = Result<()>> + Send,
     {
         self.add_entity_message_handler::<M, E, _, _>(move |subscriber, message, client, cx| {
@@ -627,8 +627,8 @@ impl Client {
     fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
     where
         M: EntityMessage,
-        E: 'static + Send + Sync,
-        H: 'static + Send + Sync + Fn(AnyHandle, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
+        E: 'static + Send,
+        H: 'static + Send + Sync + Fn(AnyModel, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
         F: 'static + Future<Output = Result<()>> + Send,
     {
         let model_type_id = TypeId::of::<E>();
@@ -666,8 +666,8 @@ impl Client {
     pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
     where
         M: EntityMessage + RequestMessage,
-        E: 'static + Send + Sync,
-        H: 'static + Send + Sync + Fn(Handle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
+        E: 'static + Send,
+        H: 'static + Send + Sync + Fn(Model<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
         F: 'static + Future<Output = Result<M::Response>> + Send,
     {
         self.add_model_message_handler(move |entity, envelope, client, cx| {
@@ -1546,7 +1546,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(
-            move |model: Handle<Model>, _: TypedEnvelope<proto::JoinProject>, _, mut cx| {
+            move |model: Model<TestModel>, _: TypedEnvelope<proto::JoinProject>, _, mut cx| {
                 match model.update(&mut cx, |model, _| model.id).unwrap() {
                     1 => done_tx1.try_send(()).unwrap(),
                     2 => done_tx2.try_send(()).unwrap(),
@@ -1555,15 +1555,15 @@ mod tests {
                 async { Ok(()) }
             },
         );
-        let model1 = cx.entity(|_| Model {
+        let model1 = cx.build_model(|_| TestModel {
             id: 1,
             subscription: None,
         });
-        let model2 = cx.entity(|_| Model {
+        let model2 = cx.build_model(|_| TestModel {
             id: 2,
             subscription: None,
         });
-        let model3 = cx.entity(|_| Model {
+        let model3 = cx.build_model(|_| TestModel {
             id: 3,
             subscription: None,
         });
@@ -1596,7 +1596,7 @@ mod tests {
         let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
         let server = FakeServer::for_client(user_id, &client, cx).await;
 
-        let model = cx.entity(|_| Model::default());
+        let model = cx.build_model(|_| TestModel::default());
         let (done_tx1, _done_rx1) = smol::channel::unbounded();
         let (done_tx2, mut done_rx2) = smol::channel::unbounded();
         let subscription1 = client.add_message_handler(
@@ -1624,11 +1624,11 @@ mod tests {
         let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
         let server = FakeServer::for_client(user_id, &client, cx).await;
 
-        let model = cx.entity(|_| Model::default());
+        let model = cx.build_model(|_| TestModel::default());
         let (done_tx, mut done_rx) = smol::channel::unbounded();
         let subscription = client.add_message_handler(
             model.clone().downgrade(),
-            move |model: Handle<Model>, _: TypedEnvelope<proto::Ping>, _, mut cx| {
+            move |model: Model<TestModel>, _: TypedEnvelope<proto::Ping>, _, mut cx| {
                 model
                     .update(&mut cx, |model, _| model.subscription.take())
                     .unwrap();
@@ -1644,7 +1644,7 @@ mod tests {
     }
 
     #[derive(Default)]
-    struct Model {
+    struct TestModel {
         id: usize,
         subscription: Option<Subscription>,
     }

crates/client2/src/telemetry.rs

@@ -5,7 +5,9 @@ use parking_lot::Mutex;
 use serde::Serialize;
 use settings2::Settings;
 use std::{env, io::Write, mem, path::PathBuf, sync::Arc, time::Duration};
-use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt};
+use sysinfo::{
+    CpuRefreshKind, Pid, PidExt, ProcessExt, ProcessRefreshKind, RefreshKind, System, SystemExt,
+};
 use tempfile::NamedTempFile;
 use util::http::HttpClient;
 use util::{channel::ReleaseChannel, TryFutureExt};
@@ -161,8 +163,16 @@ impl Telemetry {
 
         let this = self.clone();
         cx.spawn(|cx| async move {
-            let mut system = System::new_all();
-            system.refresh_all();
+            // Avoiding calling `System::new_all()`, as there have been crashes related to it
+            let refresh_kind = RefreshKind::new()
+                .with_memory() // For memory usage
+                .with_processes(ProcessRefreshKind::everything()) // For process usage
+                .with_cpu(CpuRefreshKind::everything()); // For core count
+
+            let mut system = System::new_with_specifics(refresh_kind);
+
+            // Avoiding calling `refresh_all()`, just update what we need
+            system.refresh_specifics(refresh_kind);
 
             loop {
                 // Waiting some amount of time before the first query is important to get a reasonable value
@@ -170,8 +180,7 @@ impl Telemetry {
                 const DURATION_BETWEEN_SYSTEM_EVENTS: Duration = Duration::from_secs(60);
                 smol::Timer::after(DURATION_BETWEEN_SYSTEM_EVENTS).await;
 
-                system.refresh_memory();
-                system.refresh_processes();
+                system.refresh_specifics(refresh_kind);
 
                 let current_process = Pid::from_u32(std::process::id());
                 let Some(process) = system.processes().get(&current_process) else {

crates/client2/src/test.rs

@@ -1,7 +1,7 @@
 use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
 use anyhow::{anyhow, Result};
 use futures::{stream::BoxStream, StreamExt};
-use gpui2::{Context, Executor, Handle, TestAppContext};
+use gpui2::{Context, Executor, Model, TestAppContext};
 use parking_lot::Mutex;
 use rpc2::{
     proto::{self, GetPrivateUserInfo, GetPrivateUserInfoResponse},
@@ -194,9 +194,9 @@ impl FakeServer {
         &self,
         client: Arc<Client>,
         cx: &mut TestAppContext,
-    ) -> Handle<UserStore> {
+    ) -> Model<UserStore> {
         let http_client = FakeHttpClient::with_404_response();
-        let user_store = cx.entity(|cx| UserStore::new(client, http_client, cx));
+        let user_store = cx.build_model(|cx| UserStore::new(client, http_client, cx));
         assert_eq!(
             self.receive::<proto::GetUsers>()
                 .await

crates/client2/src/user.rs

@@ -3,7 +3,7 @@ use anyhow::{anyhow, Context, Result};
 use collections::{hash_map::Entry, HashMap, HashSet};
 use feature_flags2::FeatureFlagAppExt;
 use futures::{channel::mpsc, future, AsyncReadExt, Future, StreamExt};
-use gpui2::{AsyncAppContext, EventEmitter, Handle, ImageData, ModelContext, Task};
+use gpui2::{AsyncAppContext, EventEmitter, ImageData, Model, ModelContext, Task};
 use postage::{sink::Sink, watch};
 use rpc2::proto::{RequestMessage, UsersResponse};
 use std::sync::{Arc, Weak};
@@ -122,9 +122,9 @@ impl UserStore {
         let (mut current_user_tx, current_user_rx) = watch::channel();
         let (update_contacts_tx, mut update_contacts_rx) = mpsc::unbounded();
         let rpc_subscriptions = vec![
-            client.add_message_handler(cx.weak_handle(), Self::handle_update_contacts),
-            client.add_message_handler(cx.weak_handle(), Self::handle_update_invite_info),
-            client.add_message_handler(cx.weak_handle(), Self::handle_show_contacts),
+            client.add_message_handler(cx.weak_model(), Self::handle_update_contacts),
+            client.add_message_handler(cx.weak_model(), Self::handle_update_invite_info),
+            client.add_message_handler(cx.weak_model(), Self::handle_show_contacts),
         ];
         Self {
             users: Default::default(),
@@ -213,7 +213,7 @@ impl UserStore {
     }
 
     async fn handle_update_invite_info(
-        this: Handle<Self>,
+        this: Model<Self>,
         message: TypedEnvelope<proto::UpdateInviteInfo>,
         _: Arc<Client>,
         mut cx: AsyncAppContext,
@@ -229,7 +229,7 @@ impl UserStore {
     }
 
     async fn handle_show_contacts(
-        this: Handle<Self>,
+        this: Model<Self>,
         _: TypedEnvelope<proto::ShowContacts>,
         _: Arc<Client>,
         mut cx: AsyncAppContext,
@@ -243,7 +243,7 @@ impl UserStore {
     }
 
     async fn handle_update_contacts(
-        this: Handle<Self>,
+        this: Model<Self>,
         message: TypedEnvelope<proto::UpdateContacts>,
         _: Arc<Client>,
         mut cx: AsyncAppContext,
@@ -690,7 +690,7 @@ impl User {
 impl Contact {
     async fn from_proto(
         contact: proto::Contact,
-        user_store: &Handle<UserStore>,
+        user_store: &Model<UserStore>,
         cx: &mut AsyncAppContext,
     ) -> Result<Self> {
         let user = user_store

crates/collab/Cargo.toml

@@ -3,7 +3,7 @@ authors = ["Nathan Sobo <nathan@zed.dev>"]
 default-run = "collab"
 edition = "2021"
 name = "collab"
-version = "0.24.0"
+version = "0.27.0"
 publish = false
 
 [[bin]]
@@ -73,6 +73,7 @@ git = { path = "../git", features = ["test-support"] }
 live_kit_client = { path = "../live_kit_client", features = ["test-support"] }
 lsp = { path = "../lsp", features = ["test-support"] }
 node_runtime = { path = "../node_runtime" }
+notifications = { path = "../notifications", features = ["test-support"] }
 project = { path = "../project", features = ["test-support"] }
 rpc = { path = "../rpc", features = ["test-support"] }
 settings = { path = "../settings", features = ["test-support"] }

crates/collab/migrations.sqlite/20221109000000_test_schema.sql

@@ -44,7 +44,7 @@ CREATE UNIQUE INDEX "index_rooms_on_channel_id" ON "rooms" ("channel_id");
 
 CREATE TABLE "projects" (
     "id" INTEGER PRIMARY KEY AUTOINCREMENT,
-    "room_id" INTEGER REFERENCES rooms (id) NOT NULL,
+    "room_id" INTEGER REFERENCES rooms (id) ON DELETE CASCADE NOT NULL,
     "host_user_id" INTEGER REFERENCES users (id) NOT NULL,
     "host_connection_id" INTEGER,
     "host_connection_server_id" INTEGER REFERENCES servers (id) ON DELETE CASCADE,
@@ -192,9 +192,13 @@ CREATE INDEX "index_followers_on_room_id" ON "followers" ("room_id");
 CREATE TABLE "channels" (
     "id" INTEGER PRIMARY KEY AUTOINCREMENT,
     "name" VARCHAR NOT NULL,
-    "created_at" TIMESTAMP NOT NULL DEFAULT now
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "visibility" VARCHAR NOT NULL,
+    "parent_path" TEXT
 );
 
+CREATE INDEX "index_channels_on_parent_path" ON "channels" ("parent_path");
+
 CREATE TABLE IF NOT EXISTS "channel_chat_participants" (
     "id" INTEGER PRIMARY KEY AUTOINCREMENT,
     "user_id" INTEGER NOT NULL REFERENCES users (id),
@@ -213,19 +217,22 @@ CREATE TABLE IF NOT EXISTS "channel_messages" (
     "nonce" BLOB NOT NULL
 );
 CREATE INDEX "index_channel_messages_on_channel_id" ON "channel_messages" ("channel_id");
-CREATE UNIQUE INDEX "index_channel_messages_on_nonce" ON "channel_messages" ("nonce");
+CREATE UNIQUE INDEX "index_channel_messages_on_sender_id_nonce" ON "channel_messages" ("sender_id", "nonce");
 
-CREATE TABLE "channel_paths" (
-    "id_path" TEXT NOT NULL PRIMARY KEY,
-    "channel_id" INTEGER NOT NULL REFERENCES channels (id) ON DELETE CASCADE
+CREATE TABLE "channel_message_mentions" (
+    "message_id" INTEGER NOT NULL REFERENCES channel_messages (id) ON DELETE CASCADE,
+    "start_offset" INTEGER NOT NULL,
+    "end_offset" INTEGER NOT NULL,
+    "user_id" INTEGER NOT NULL REFERENCES users (id) ON DELETE CASCADE,
+    PRIMARY KEY(message_id, start_offset)
 );
-CREATE INDEX "index_channel_paths_on_channel_id" ON "channel_paths" ("channel_id");
 
 CREATE TABLE "channel_members" (
     "id" INTEGER PRIMARY KEY AUTOINCREMENT,
     "channel_id" INTEGER NOT NULL REFERENCES channels (id) ON DELETE CASCADE,
     "user_id" INTEGER NOT NULL REFERENCES users (id) ON DELETE CASCADE,
     "admin" BOOLEAN NOT NULL DEFAULT false,
+    "role" VARCHAR,
     "accepted" BOOLEAN NOT NULL DEFAULT false,
     "updated_at" TIMESTAMP NOT NULL DEFAULT now
 );
@@ -312,3 +319,26 @@ CREATE TABLE IF NOT EXISTS "observed_channel_messages" (
 );
 
 CREATE UNIQUE INDEX "index_observed_channel_messages_user_and_channel_id" ON "observed_channel_messages" ("user_id", "channel_id");
+
+CREATE TABLE "notification_kinds" (
+    "id" INTEGER PRIMARY KEY AUTOINCREMENT,
+    "name" VARCHAR NOT NULL
+);
+
+CREATE UNIQUE INDEX "index_notification_kinds_on_name" ON "notification_kinds" ("name");
+
+CREATE TABLE "notifications" (
+    "id" INTEGER PRIMARY KEY AUTOINCREMENT,
+    "created_at" TIMESTAMP NOT NULL default CURRENT_TIMESTAMP,
+    "recipient_id" INTEGER NOT NULL REFERENCES users (id) ON DELETE CASCADE,
+    "kind" INTEGER NOT NULL REFERENCES notification_kinds (id),
+    "entity_id" INTEGER,
+    "content" TEXT,
+    "is_read" BOOLEAN NOT NULL DEFAULT FALSE,
+    "response" BOOLEAN
+);
+
+CREATE INDEX
+    "index_notifications_on_recipient_id_is_read_kind_entity_id"
+    ON "notifications"
+    ("recipient_id", "is_read", "kind", "entity_id");

crates/collab/migrations/20231004130100_create_notifications.sql

@@ -0,0 +1,22 @@
+CREATE TABLE "notification_kinds" (
+    "id" SERIAL PRIMARY KEY,
+    "name" VARCHAR NOT NULL
+);
+
+CREATE UNIQUE INDEX "index_notification_kinds_on_name" ON "notification_kinds" ("name");
+
+CREATE TABLE notifications (
+    "id" SERIAL PRIMARY KEY,
+    "created_at" TIMESTAMP NOT NULL DEFAULT now(),
+    "recipient_id" INTEGER NOT NULL REFERENCES users (id) ON DELETE CASCADE,
+    "kind" INTEGER NOT NULL REFERENCES notification_kinds (id),
+    "entity_id" INTEGER,
+    "content" TEXT,
+    "is_read" BOOLEAN NOT NULL DEFAULT FALSE,
+    "response" BOOLEAN
+);
+
+CREATE INDEX
+    "index_notifications_on_recipient_id_is_read_kind_entity_id"
+    ON "notifications"
+    ("recipient_id", "is_read", "kind", "entity_id");

crates/collab/migrations/20231011214412_add_guest_role.sql

@@ -0,0 +1,4 @@
+ALTER TABLE channel_members ADD COLUMN role TEXT;
+UPDATE channel_members SET role = CASE WHEN admin THEN 'admin' ELSE 'member' END;
+
+ALTER TABLE channels ADD COLUMN visibility TEXT NOT NULL DEFAULT 'members';

crates/collab/migrations/20231017185833_projects_room_id_fkey_on_delete_cascade.sql

@@ -0,0 +1,8 @@
+-- Add migration script here
+
+ALTER TABLE projects
+  DROP CONSTRAINT projects_room_id_fkey,
+  ADD CONSTRAINT  projects_room_id_fkey
+    FOREIGN KEY (room_id)
+    REFERENCES rooms (id)
+    ON DELETE CASCADE;

crates/collab/migrations/20231018102700_create_mentions.sql

@@ -0,0 +1,11 @@
+CREATE TABLE "channel_message_mentions" (
+    "message_id" INTEGER NOT NULL REFERENCES channel_messages (id) ON DELETE CASCADE,
+    "start_offset" INTEGER NOT NULL,
+    "end_offset" INTEGER NOT NULL,
+    "user_id" INTEGER NOT NULL REFERENCES users (id) ON DELETE CASCADE,
+    PRIMARY KEY(message_id, start_offset)
+);
+
+-- We use 'on conflict update' with this index, so it should be per-user.
+CREATE UNIQUE INDEX "index_channel_messages_on_sender_id_nonce" ON "channel_messages" ("sender_id", "nonce");
+DROP INDEX "index_channel_messages_on_nonce";

crates/collab/migrations/20231024085546_move_channel_paths_to_channels_table.sql

@@ -0,0 +1,12 @@
+ALTER TABLE channels ADD COLUMN parent_path TEXT;
+
+UPDATE channels
+SET parent_path = substr(
+    channel_paths.id_path,
+    2,
+    length(channel_paths.id_path) - length('/' || channel_paths.channel_id::text || '/')
+)
+FROM channel_paths
+WHERE channel_paths.channel_id = channels.id;
+
+CREATE INDEX "index_channels_on_parent_path" ON "channels" ("parent_path");

crates/collab/src/bin/seed.rs

@@ -71,7 +71,6 @@ async fn main() {
                     db::NewUserParams {
                         github_login: github_user.login,
                         github_user_id: github_user.id,
-                        invite_count: 5,
                     },
                 )
                 .await

crates/collab/src/db.rs

@@ -20,7 +20,7 @@ use rpc::{
 };
 use sea_orm::{
     entity::prelude::*,
-    sea_query::{Alias, Expr, OnConflict, Query},
+    sea_query::{Alias, Expr, OnConflict},
     ActiveValue, Condition, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbErr,
     FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect, Statement,
     TransactionTrait,
@@ -47,14 +47,14 @@ pub use ids::*;
 pub use sea_orm::ConnectOptions;
 pub use tables::user::Model as User;
 
-use self::queries::channels::ChannelGraph;
-
 pub struct Database {
     options: ConnectOptions,
     pool: DatabaseConnection,
     rooms: DashMap<RoomId, Arc<Mutex<()>>>,
     rng: Mutex<StdRng>,
     executor: Executor,
+    notification_kinds_by_id: HashMap<NotificationKindId, &'static str>,
+    notification_kinds_by_name: HashMap<String, NotificationKindId>,
     #[cfg(test)]
     runtime: Option<tokio::runtime::Runtime>,
 }
@@ -69,6 +69,8 @@ impl Database {
             pool: sea_orm::Database::connect(options).await?,
             rooms: DashMap::with_capacity(16384),
             rng: Mutex::new(StdRng::seed_from_u64(0)),
+            notification_kinds_by_id: HashMap::default(),
+            notification_kinds_by_name: HashMap::default(),
             executor,
             #[cfg(test)]
             runtime: None,
@@ -121,6 +123,11 @@ impl Database {
         Ok(new_migrations)
     }
 
+    pub async fn initialize_static_data(&mut self) -> Result<()> {
+        self.initialize_notification_kinds().await?;
+        Ok(())
+    }
+
     pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
     where
         F: Send + Fn(TransactionHandle) -> Fut,
@@ -361,18 +368,9 @@ impl<T> RoomGuard<T> {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum Contact {
-    Accepted {
-        user_id: UserId,
-        should_notify: bool,
-        busy: bool,
-    },
-    Outgoing {
-        user_id: UserId,
-    },
-    Incoming {
-        user_id: UserId,
-        should_notify: bool,
-    },
+    Accepted { user_id: UserId, busy: bool },
+    Outgoing { user_id: UserId },
+    Incoming { user_id: UserId },
 }
 
 impl Contact {
@@ -385,6 +383,15 @@ impl Contact {
     }
 }
 
+pub type NotificationBatch = Vec<(UserId, proto::Notification)>;
+
+pub struct CreatedChannelMessage {
+    pub message_id: MessageId,
+    pub participant_connection_ids: Vec<ConnectionId>,
+    pub channel_members: Vec<UserId>,
+    pub notifications: NotificationBatch,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]
 pub struct Invite {
     pub email_address: String,
@@ -417,7 +424,6 @@ pub struct WaitlistSummary {
 pub struct NewUserParams {
     pub github_login: String,
     pub github_user_id: i32,
-    pub invite_count: i32,
 }
 
 #[derive(Debug)]
@@ -428,17 +434,115 @@ pub struct NewUserResult {
     pub signup_device_id: Option<String>,
 }
 
-#[derive(FromQueryResult, Debug, PartialEq, Eq, Hash)]
+#[derive(Debug)]
+pub struct MoveChannelResult {
+    pub participants_to_update: HashMap<UserId, ChannelsForUser>,
+    pub participants_to_remove: HashSet<UserId>,
+    pub moved_channels: HashSet<ChannelId>,
+}
+
+#[derive(Debug)]
+pub struct RenameChannelResult {
+    pub channel: Channel,
+    pub participants_to_update: HashMap<UserId, Channel>,
+}
+
+#[derive(Debug)]
+pub struct CreateChannelResult {
+    pub channel: Channel,
+    pub participants_to_update: Vec<(UserId, ChannelsForUser)>,
+}
+
+#[derive(Debug)]
+pub struct SetChannelVisibilityResult {
+    pub participants_to_update: HashMap<UserId, ChannelsForUser>,
+    pub participants_to_remove: HashSet<UserId>,
+    pub channels_to_remove: Vec<ChannelId>,
+}
+
+#[derive(Debug)]
+pub struct MembershipUpdated {
+    pub channel_id: ChannelId,
+    pub new_channels: ChannelsForUser,
+    pub removed_channels: Vec<ChannelId>,
+}
+
+#[derive(Debug)]
+pub enum SetMemberRoleResult {
+    InviteUpdated(Channel),
+    MembershipUpdated(MembershipUpdated),
+}
+
+#[derive(Debug)]
+pub struct InviteMemberResult {
+    pub channel: Channel,
+    pub notifications: NotificationBatch,
+}
+
+#[derive(Debug)]
+pub struct RespondToChannelInvite {
+    pub membership_update: Option<MembershipUpdated>,
+    pub notifications: NotificationBatch,
+}
+
+#[derive(Debug)]
+pub struct RemoveChannelMemberResult {
+    pub membership_update: MembershipUpdated,
+    pub notification_id: Option<NotificationId>,
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Channel {
     pub id: ChannelId,
     pub name: String,
+    pub visibility: ChannelVisibility,
+    pub role: ChannelRole,
+    pub parent_path: Vec<ChannelId>,
+}
+
+impl Channel {
+    fn from_model(value: channel::Model, role: ChannelRole) -> Self {
+        Channel {
+            id: value.id,
+            visibility: value.visibility,
+            name: value.clone().name,
+            role,
+            parent_path: value.ancestors().collect(),
+        }
+    }
+
+    pub fn to_proto(&self) -> proto::Channel {
+        proto::Channel {
+            id: self.id.to_proto(),
+            name: self.name.clone(),
+            visibility: self.visibility.into(),
+            role: self.role.into(),
+            parent_path: self.parent_path.iter().map(|c| c.to_proto()).collect(),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ChannelMember {
+    pub role: ChannelRole,
+    pub user_id: UserId,
+    pub kind: proto::channel_member::Kind,
+}
+
+impl ChannelMember {
+    pub fn to_proto(&self) -> proto::ChannelMember {
+        proto::ChannelMember {
+            role: self.role.into(),
+            user_id: self.user_id.to_proto(),
+            kind: self.kind.into(),
+        }
+    }
 }
 
 #[derive(Debug, PartialEq)]
 pub struct ChannelsForUser {
-    pub channels: ChannelGraph,
+    pub channels: Vec<Channel>,
     pub channel_participants: HashMap<ChannelId, Vec<UserId>>,
-    pub channels_with_admin_privileges: HashSet<ChannelId>,
     pub unseen_buffer_changes: Vec<proto::UnseenChannelBufferChange>,
     pub channel_messages: Vec<proto::UnseenChannelMessage>,
 }

crates/collab/src/db/ids.rs

@@ -1,4 +1,5 @@
 use crate::Result;
+use rpc::proto;
 use sea_orm::{entity::prelude::*, DbErr};
 use serde::{Deserialize, Serialize};
 
@@ -80,3 +81,119 @@ id_type!(SignupId);
 id_type!(UserId);
 id_type!(ChannelBufferCollaboratorId);
 id_type!(FlagId);
+id_type!(NotificationId);
+id_type!(NotificationKindId);
+
+#[derive(Eq, PartialEq, Copy, Clone, Debug, EnumIter, DeriveActiveEnum, Default, Hash)]
+#[sea_orm(rs_type = "String", db_type = "String(None)")]
+pub enum ChannelRole {
+    #[sea_orm(string_value = "admin")]
+    Admin,
+    #[sea_orm(string_value = "member")]
+    #[default]
+    Member,
+    #[sea_orm(string_value = "guest")]
+    Guest,
+    #[sea_orm(string_value = "banned")]
+    Banned,
+}
+
+impl ChannelRole {
+    pub fn should_override(&self, other: Self) -> bool {
+        use ChannelRole::*;
+        match self {
+            Admin => matches!(other, Member | Banned | Guest),
+            Member => matches!(other, Banned | Guest),
+            Banned => matches!(other, Guest),
+            Guest => false,
+        }
+    }
+
+    pub fn max(&self, other: Self) -> Self {
+        if self.should_override(other) {
+            *self
+        } else {
+            other
+        }
+    }
+
+    pub fn can_see_all_descendants(&self) -> bool {
+        use ChannelRole::*;
+        match self {
+            Admin | Member => true,
+            Guest | Banned => false,
+        }
+    }
+
+    pub fn can_only_see_public_descendants(&self) -> bool {
+        use ChannelRole::*;
+        match self {
+            Guest => true,
+            Admin | Member | Banned => false,
+        }
+    }
+}
+
+impl From<proto::ChannelRole> for ChannelRole {
+    fn from(value: proto::ChannelRole) -> Self {
+        match value {
+            proto::ChannelRole::Admin => ChannelRole::Admin,
+            proto::ChannelRole::Member => ChannelRole::Member,
+            proto::ChannelRole::Guest => ChannelRole::Guest,
+            proto::ChannelRole::Banned => ChannelRole::Banned,
+        }
+    }
+}
+
+impl Into<proto::ChannelRole> for ChannelRole {
+    fn into(self) -> proto::ChannelRole {
+        match self {
+            ChannelRole::Admin => proto::ChannelRole::Admin,
+            ChannelRole::Member => proto::ChannelRole::Member,
+            ChannelRole::Guest => proto::ChannelRole::Guest,
+            ChannelRole::Banned => proto::ChannelRole::Banned,
+        }
+    }
+}
+
+impl Into<i32> for ChannelRole {
+    fn into(self) -> i32 {
+        let proto: proto::ChannelRole = self.into();
+        proto.into()
+    }
+}
+
+#[derive(Eq, PartialEq, Copy, Clone, Debug, EnumIter, DeriveActiveEnum, Default, Hash)]
+#[sea_orm(rs_type = "String", db_type = "String(None)")]
+pub enum ChannelVisibility {
+    #[sea_orm(string_value = "public")]
+    Public,
+    #[sea_orm(string_value = "members")]
+    #[default]
+    Members,
+}
+
+impl From<proto::ChannelVisibility> for ChannelVisibility {
+    fn from(value: proto::ChannelVisibility) -> Self {
+        match value {
+            proto::ChannelVisibility::Public => ChannelVisibility::Public,
+            proto::ChannelVisibility::Members => ChannelVisibility::Members,
+        }
+    }
+}
+
+impl Into<proto::ChannelVisibility> for ChannelVisibility {
+    fn into(self) -> proto::ChannelVisibility {
+        match self {
+            ChannelVisibility::Public => proto::ChannelVisibility::Public,
+            ChannelVisibility::Members => proto::ChannelVisibility::Members,
+        }
+    }
+}
+
+impl Into<i32> for ChannelVisibility {
+    fn into(self) -> i32 {
+        let proto: proto::ChannelVisibility = self.into();
+        proto.into()
+    }
+}

crates/collab/src/db/queries.rs

@@ -5,6 +5,7 @@ pub mod buffers;
 pub mod channels;
 pub mod contacts;
 pub mod messages;
+pub mod notifications;
 pub mod projects;
 pub mod rooms;
 pub mod servers;

crates/collab/src/db/queries/access_tokens.rs

@@ -1,4 +1,5 @@
 use super::*;
+use sea_orm::sea_query::Query;
 
 impl Database {
     pub async fn create_access_token(

crates/collab/src/db/queries/buffers.rs

@@ -16,7 +16,8 @@ impl Database {
         connection: ConnectionId,
     ) -> Result<proto::JoinChannelBufferResponse> {
         self.transaction(|tx| async move {
-            self.check_user_is_channel_member(channel_id, user_id, &tx)
+            let channel = self.get_channel_internal(channel_id, &*tx).await?;
+            self.check_user_is_channel_participant(&channel, user_id, &tx)
                 .await?;
 
             let buffer = channel::Model {
@@ -129,9 +130,11 @@ impl Database {
         self.transaction(|tx| async move {
             let mut results = Vec::new();
             for client_buffer in buffers {
-                let channel_id = ChannelId::from_proto(client_buffer.channel_id);
+                let channel = self
+                    .get_channel_internal(ChannelId::from_proto(client_buffer.channel_id), &*tx)
+                    .await?;
                 if self
-                    .check_user_is_channel_member(channel_id, user_id, &*tx)
+                    .check_user_is_channel_participant(&channel, user_id, &*tx)
                     .await
                     .is_err()
                 {
@@ -139,9 +142,9 @@ impl Database {
                     continue;
                 }
 
-                let buffer = self.get_channel_buffer(channel_id, &*tx).await?;
+                let buffer = self.get_channel_buffer(channel.id, &*tx).await?;
                 let mut collaborators = channel_buffer_collaborator::Entity::find()
-                    .filter(channel_buffer_collaborator::Column::ChannelId.eq(channel_id))
+                    .filter(channel_buffer_collaborator::Column::ChannelId.eq(channel.id))
                     .all(&*tx)
                     .await?;
 
@@ -439,7 +442,8 @@ impl Database {
         Vec<proto::VectorClockEntry>,
     )> {
         self.transaction(move |tx| async move {
-            self.check_user_is_channel_member(channel_id, user, &*tx)
+            let channel = self.get_channel_internal(channel_id, &*tx).await?;
+            self.check_user_is_channel_member(&channel, user, &*tx)
                 .await?;
 
             let buffer = buffer::Entity::find()
@@ -482,7 +486,7 @@ impl Database {
                 )
                 .await?;
 
-                channel_members = self.get_channel_members_internal(channel_id, &*tx).await?;
+                channel_members = self.get_channel_participants(&channel, &*tx).await?;
                 let collaborators = self
                     .get_channel_buffer_collaborators_internal(channel_id, &*tx)
                     .await?;

crates/collab/src/db/queries/contacts.rs

@@ -8,7 +8,6 @@ impl Database {
             user_id_b: UserId,
             a_to_b: bool,
             accepted: bool,
-            should_notify: bool,
             user_a_busy: bool,
             user_b_busy: bool,
         }
@@ -53,7 +52,6 @@ impl Database {
                     if db_contact.accepted {
                         contacts.push(Contact::Accepted {
                             user_id: db_contact.user_id_b,
-                            should_notify: db_contact.should_notify && db_contact.a_to_b,
                             busy: db_contact.user_b_busy,
                         });
                     } else if db_contact.a_to_b {
@@ -63,19 +61,16 @@ impl Database {
                     } else {
                         contacts.push(Contact::Incoming {
                             user_id: db_contact.user_id_b,
-                            should_notify: db_contact.should_notify,
                         });
                     }
                 } else if db_contact.accepted {
                     contacts.push(Contact::Accepted {
                         user_id: db_contact.user_id_a,
-                        should_notify: db_contact.should_notify && !db_contact.a_to_b,
                         busy: db_contact.user_a_busy,
                     });
                 } else if db_contact.a_to_b {
                     contacts.push(Contact::Incoming {
                         user_id: db_contact.user_id_a,
-                        should_notify: db_contact.should_notify,
                     });
                 } else {
                     contacts.push(Contact::Outgoing {
@@ -124,7 +119,11 @@ impl Database {
         .await
     }
 
-    pub async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> {
+    pub async fn send_contact_request(
+        &self,
+        sender_id: UserId,
+        receiver_id: UserId,
+    ) -> Result<NotificationBatch> {
         self.transaction(|tx| async move {
             let (id_a, id_b, a_to_b) = if sender_id < receiver_id {
                 (sender_id, receiver_id, true)
@@ -161,11 +160,22 @@ impl Database {
             .exec_without_returning(&*tx)
             .await?;
 
-            if rows_affected == 1 {
-                Ok(())
-            } else {
-                Err(anyhow!("contact already requested"))?
+            if rows_affected == 0 {
+                Err(anyhow!("contact already requested"))?;
             }
+
+            Ok(self
+                .create_notification(
+                    receiver_id,
+                    rpc::Notification::ContactRequest {
+                        sender_id: sender_id.to_proto(),
+                    },
+                    true,
+                    &*tx,
+                )
+                .await?
+                .into_iter()
+                .collect())
         })
         .await
     }
@@ -179,7 +189,11 @@ impl Database {
     ///
     /// * `requester_id` - The user that initiates this request
     /// * `responder_id` - The user that will be removed
-    pub async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<bool> {
+    pub async fn remove_contact(
+        &self,
+        requester_id: UserId,
+        responder_id: UserId,
+    ) -> Result<(bool, Option<NotificationId>)> {
         self.transaction(|tx| async move {
             let (id_a, id_b) = if responder_id < requester_id {
                 (responder_id, requester_id)
@@ -198,7 +212,21 @@ impl Database {
                 .ok_or_else(|| anyhow!("no such contact"))?;
 
             contact::Entity::delete_by_id(contact.id).exec(&*tx).await?;
-            Ok(contact.accepted)
+
+            let mut deleted_notification_id = None;
+            if !contact.accepted {
+                deleted_notification_id = self
+                    .remove_notification(
+                        responder_id,
+                        rpc::Notification::ContactRequest {
+                            sender_id: requester_id.to_proto(),
+                        },
+                        &*tx,
+                    )
+                    .await?;
+            }
+
+            Ok((contact.accepted, deleted_notification_id))
         })
         .await
     }
@@ -249,7 +277,7 @@ impl Database {
         responder_id: UserId,
         requester_id: UserId,
         accept: bool,
-    ) -> Result<()> {
+    ) -> Result<NotificationBatch> {
         self.transaction(|tx| async move {
             let (id_a, id_b, a_to_b) = if responder_id < requester_id {
                 (responder_id, requester_id, false)
@@ -287,11 +315,38 @@ impl Database {
                 result.rows_affected
             };
 
-            if rows_affected == 1 {
-                Ok(())
-            } else {
+            if rows_affected == 0 {
                 Err(anyhow!("no such contact request"))?
             }
+
+            let mut notifications = Vec::new();
+            notifications.extend(
+                self.mark_notification_as_read_with_response(
+                    responder_id,
+                    &rpc::Notification::ContactRequest {
+                        sender_id: requester_id.to_proto(),
+                    },
+                    accept,
+                    &*tx,
+                )
+                .await?,
+            );
+
+            if accept {
+                notifications.extend(
+                    self.create_notification(
+                        requester_id,
+                        rpc::Notification::ContactRequestAccepted {
+                            responder_id: responder_id.to_proto(),
+                        },
+                        true,
+                        &*tx,
+                    )
+                    .await?,
+                );
+            }
+
+            Ok(notifications)
         })
         .await
     }

crates/collab/src/db/queries/messages.rs

@@ -1,4 +1,6 @@
 use super::*;
+use rpc::Notification;
+use sea_orm::TryInsertResult;
 use time::OffsetDateTime;
 
 impl Database {
@@ -9,7 +11,8 @@ impl Database {
         user_id: UserId,
     ) -> Result<()> {
         self.transaction(|tx| async move {
-            self.check_user_is_channel_member(channel_id, user_id, &*tx)
+            let channel = self.get_channel_internal(channel_id, &*tx).await?;
+            self.check_user_is_channel_participant(&channel, user_id, &*tx)
                 .await?;
             channel_chat_participant::ActiveModel {
                 id: ActiveValue::NotSet,
@@ -77,7 +80,8 @@ impl Database {
         before_message_id: Option<MessageId>,
     ) -> Result<Vec<proto::ChannelMessage>> {
         self.transaction(|tx| async move {
-            self.check_user_is_channel_member(channel_id, user_id, &*tx)
+            let channel = self.get_channel_internal(channel_id, &*tx).await?;
+            self.check_user_is_channel_participant(&channel, user_id, &*tx)
                 .await?;
 
             let mut condition =
@@ -87,33 +91,103 @@ impl Database {
                 condition = condition.add(channel_message::Column::Id.lt(before_message_id));
             }
 
-            let mut rows = channel_message::Entity::find()
+            let rows = channel_message::Entity::find()
                 .filter(condition)
                 .order_by_desc(channel_message::Column::Id)
                 .limit(count as u64)
-                .stream(&*tx)
+                .all(&*tx)
                 .await?;
 
-            let mut messages = Vec::new();
-            while let Some(row) = rows.next().await {
-                let row = row?;
+            self.load_channel_messages(rows, &*tx).await
+        })
+        .await
+    }
+
+    pub async fn get_channel_messages_by_id(
+        &self,
+        user_id: UserId,
+        message_ids: &[MessageId],
+    ) -> Result<Vec<proto::ChannelMessage>> {
+        self.transaction(|tx| async move {
+            let rows = channel_message::Entity::find()
+                .filter(channel_message::Column::Id.is_in(message_ids.iter().copied()))
+                .order_by_desc(channel_message::Column::Id)
+                .all(&*tx)
+                .await?;
+
+            let mut channels = HashMap::<ChannelId, channel::Model>::default();
+            for row in &rows {
+                channels.insert(
+                    row.channel_id,
+                    self.get_channel_internal(row.channel_id, &*tx).await?,
+                );
+            }
+
+            for (_, channel) in channels {
+                self.check_user_is_channel_participant(&channel, user_id, &*tx)
+                    .await?;
+            }
+
+            let messages = self.load_channel_messages(rows, &*tx).await?;
+            Ok(messages)
+        })
+        .await
+    }
+
+    async fn load_channel_messages(
+        &self,
+        rows: Vec<channel_message::Model>,
+        tx: &DatabaseTransaction,
+    ) -> Result<Vec<proto::ChannelMessage>> {
+        let mut messages = rows
+            .into_iter()
+            .map(|row| {
                 let nonce = row.nonce.as_u64_pair();
-                messages.push(proto::ChannelMessage {
+                proto::ChannelMessage {
                     id: row.id.to_proto(),
                     sender_id: row.sender_id.to_proto(),
                     body: row.body,
                     timestamp: row.sent_at.assume_utc().unix_timestamp() as u64,
+                    mentions: vec![],
                     nonce: Some(proto::Nonce {
                         upper_half: nonce.0,
                         lower_half: nonce.1,
                     }),
-                });
+                }
+            })
+            .collect::<Vec<_>>();
+        messages.reverse();
+
+        let mut mentions = channel_message_mention::Entity::find()
+            .filter(channel_message_mention::Column::MessageId.is_in(messages.iter().map(|m| m.id)))
+            .order_by_asc(channel_message_mention::Column::MessageId)
+            .order_by_asc(channel_message_mention::Column::StartOffset)
+            .stream(&*tx)
+            .await?;
+
+        let mut message_ix = 0;
+        while let Some(mention) = mentions.next().await {
+            let mention = mention?;
+            let message_id = mention.message_id.to_proto();
+            while let Some(message) = messages.get_mut(message_ix) {
+                if message.id < message_id {
+                    message_ix += 1;
+                } else {
+                    if message.id == message_id {
+                        message.mentions.push(proto::ChatMention {
+                            range: Some(proto::Range {
+                                start: mention.start_offset as u64,
+                                end: mention.end_offset as u64,
+                            }),
+                            user_id: mention.user_id.to_proto(),
+                        });
+                    }
+                    break;
+                }
             }
-            drop(rows);
-            messages.reverse();
-            Ok(messages)
-        })
-        .await
+        }
+
+        Ok(messages)
     }
 
     pub async fn create_channel_message(
@@ -121,10 +195,15 @@ impl Database {
         channel_id: ChannelId,
         user_id: UserId,
         body: &str,
+        mentions: &[proto::ChatMention],
         timestamp: OffsetDateTime,
         nonce: u128,
-    ) -> Result<(MessageId, Vec<ConnectionId>, Vec<UserId>)> {
+    ) -> Result<CreatedChannelMessage> {
         self.transaction(|tx| async move {
+            let channel = self.get_channel_internal(channel_id, &*tx).await?;
+            self.check_user_is_channel_participant(&channel, user_id, &*tx)
+                .await?;
+
             let mut rows = channel_chat_participant::Entity::find()
                 .filter(channel_chat_participant::Column::ChannelId.eq(channel_id))
                 .stream(&*tx)
@@ -150,7 +229,7 @@ impl Database {
             let timestamp = timestamp.to_offset(time::UtcOffset::UTC);
             let timestamp = time::PrimitiveDateTime::new(timestamp.date(), timestamp.time());
 
-            let message = channel_message::Entity::insert(channel_message::ActiveModel {
+            let result = channel_message::Entity::insert(channel_message::ActiveModel {
                 channel_id: ActiveValue::Set(channel_id),
                 sender_id: ActiveValue::Set(user_id),
                 body: ActiveValue::Set(body.to_string()),
@@ -159,35 +238,85 @@ impl Database {
                 id: ActiveValue::NotSet,
             })
             .on_conflict(
-                OnConflict::column(channel_message::Column::Nonce)
-                    .update_column(channel_message::Column::Nonce)
-                    .to_owned(),
+                OnConflict::columns([
+                    channel_message::Column::SenderId,
+                    channel_message::Column::Nonce,
+                ])
+                .do_nothing()
+                .to_owned(),
             )
+            .do_nothing()
             .exec(&*tx)
             .await?;
 
-            #[derive(Debug, Clone, Copy, EnumIter, DeriveColumn)]
-            enum QueryConnectionId {
-                ConnectionId,
+            let message_id;
+            let mut notifications = Vec::new();
+            match result {
+                TryInsertResult::Inserted(result) => {
+                    message_id = result.last_insert_id;
+                    let mentioned_user_ids =
+                        mentions.iter().map(|m| m.user_id).collect::<HashSet<_>>();
+                    let mentions = mentions
+                        .iter()
+                        .filter_map(|mention| {
+                            let range = mention.range.as_ref()?;
+                            if !body.is_char_boundary(range.start as usize)
+                                || !body.is_char_boundary(range.end as usize)
+                            {
+                                return None;
+                            }
+                            Some(channel_message_mention::ActiveModel {
+                                message_id: ActiveValue::Set(message_id),
+                                start_offset: ActiveValue::Set(range.start as i32),
+                                end_offset: ActiveValue::Set(range.end as i32),
+                                user_id: ActiveValue::Set(UserId::from_proto(mention.user_id)),
+                            })
+                        })
+                        .collect::<Vec<_>>();
+                    if !mentions.is_empty() {
+                        channel_message_mention::Entity::insert_many(mentions)
+                            .exec(&*tx)
+                            .await?;
+                    }
+
+                    for mentioned_user in mentioned_user_ids {
+                        notifications.extend(
+                            self.create_notification(
+                                UserId::from_proto(mentioned_user),
+                                rpc::Notification::ChannelMessageMention {
+                                    message_id: message_id.to_proto(),
+                                    sender_id: user_id.to_proto(),
+                                    channel_id: channel_id.to_proto(),
+                                },
+                                false,
+                                &*tx,
+                            )
+                            .await?,
+                        );
+                    }
+
+                    self.observe_channel_message_internal(channel_id, user_id, message_id, &*tx)
+                        .await?;
+                }
+                _ => {
+                    message_id = channel_message::Entity::find()
+                        .filter(channel_message::Column::Nonce.eq(Uuid::from_u128(nonce)))
+                        .one(&*tx)
+                        .await?
+                        .ok_or_else(|| anyhow!("failed to insert message"))?
+                        .id;
+                }
             }
 
-            // Observe this message for the sender
-            self.observe_channel_message_internal(
-                channel_id,
-                user_id,
-                message.last_insert_id,
-                &*tx,
-            )
-            .await?;
-
-            let mut channel_members = self.get_channel_members_internal(channel_id, &*tx).await?;
+            let mut channel_members = self.get_channel_participants(&channel, &*tx).await?;
             channel_members.retain(|member| !participant_user_ids.contains(member));
 
-            Ok((
-                message.last_insert_id,
+            Ok(CreatedChannelMessage {
+                message_id,
                 participant_connection_ids,
                 channel_members,
-            ))
+                notifications,
+            })
         })
         .await
     }
@@ -197,11 +326,24 @@ impl Database {
         channel_id: ChannelId,
         user_id: UserId,
         message_id: MessageId,
-    ) -> Result<()> {
+    ) -> Result<NotificationBatch> {
         self.transaction(|tx| async move {
             self.observe_channel_message_internal(channel_id, user_id, message_id, &*tx)
                 .await?;
-            Ok(())
+            let mut batch = NotificationBatch::default();
+            batch.extend(
+                self.mark_notification_as_read(
+                    user_id,
+                    &Notification::ChannelMessageMention {
+                        message_id: message_id.to_proto(),
+                        sender_id: Default::default(),
+                        channel_id: Default::default(),
+                    },
+                    &*tx,
+                )
+                .await?,
+            );
+            Ok(batch)
         })
         .await
     }
@@ -337,8 +479,23 @@ impl Database {
                 .filter(channel_message::Column::SenderId.eq(user_id))
                 .exec(&*tx)
                 .await?;
+
             if result.rows_affected == 0 {
-                Err(anyhow!("no such message"))?;
+                let channel = self.get_channel_internal(channel_id, &*tx).await?;
+                if self
+                    .check_user_is_channel_admin(&channel, user_id, &*tx)
+                    .await
+                    .is_ok()
+                {
+                    let result = channel_message::Entity::delete_by_id(message_id)
+                        .exec(&*tx)
+                        .await?;
+                    if result.rows_affected == 0 {
+                        Err(anyhow!("no such message"))?;
+                    }
+                } else {
+                    Err(anyhow!("operation could not be completed"))?;
+                }
             }
 
             Ok(participant_connection_ids)

crates/collab/src/db/queries/notifications.rs

@@ -0,0 +1,262 @@
+use super::*;
+use rpc::Notification;
+
+impl Database {
+    pub async fn initialize_notification_kinds(&mut self) -> Result<()> {
+        notification_kind::Entity::insert_many(Notification::all_variant_names().iter().map(
+            |kind| notification_kind::ActiveModel {
+                name: ActiveValue::Set(kind.to_string()),
+                ..Default::default()
+            },
+        ))
+        .on_conflict(OnConflict::new().do_nothing().to_owned())
+        .exec_without_returning(&self.pool)
+        .await?;
+
+        let mut rows = notification_kind::Entity::find().stream(&self.pool).await?;
+        while let Some(row) = rows.next().await {
+            let row = row?;
+            self.notification_kinds_by_name.insert(row.name, row.id);
+        }
+
+        for name in Notification::all_variant_names() {
+            if let Some(id) = self.notification_kinds_by_name.get(*name).copied() {
+                self.notification_kinds_by_id.insert(id, name);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn get_notifications(
+        &self,
+        recipient_id: UserId,
+        limit: usize,
+        before_id: Option<NotificationId>,
+    ) -> Result<Vec<proto::Notification>> {
+        self.transaction(|tx| async move {
+            let mut result = Vec::new();
+            let mut condition =
+                Condition::all().add(notification::Column::RecipientId.eq(recipient_id));
+
+            if let Some(before_id) = before_id {
+                condition = condition.add(notification::Column::Id.lt(before_id));
+            }
+
+            let mut rows = notification::Entity::find()
+                .filter(condition)
+                .order_by_desc(notification::Column::Id)
+                .limit(limit as u64)
+                .stream(&*tx)
+                .await?;
+            while let Some(row) = rows.next().await {
+                let row = row?;
+                let kind = row.kind;
+                if let Some(proto) = model_to_proto(self, row) {
+                    result.push(proto);
+                } else {
+                    log::warn!("unknown notification kind {:?}", kind);
+                }
+            }
+            result.reverse();
+            Ok(result)
+        })
+        .await
+    }
+
+    /// Create a notification. If `avoid_duplicates` is set to true, then avoid
+    /// creating a new notification if the given recipient already has an
+    /// unread notification with the given kind and entity id.
+    pub async fn create_notification(
+        &self,
+        recipient_id: UserId,
+        notification: Notification,
+        avoid_duplicates: bool,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<(UserId, proto::Notification)>> {
+        if avoid_duplicates {
+            if self
+                .find_notification(recipient_id, &notification, tx)
+                .await?
+                .is_some()
+            {
+                return Ok(None);
+            }
+        }
+
+        let proto = notification.to_proto();
+        let kind = notification_kind_from_proto(self, &proto)?;
+        let model = notification::ActiveModel {
+            recipient_id: ActiveValue::Set(recipient_id),
+            kind: ActiveValue::Set(kind),
+            entity_id: ActiveValue::Set(proto.entity_id.map(|id| id as i32)),
+            content: ActiveValue::Set(proto.content.clone()),
+            ..Default::default()
+        }
+        .save(&*tx)
+        .await?;
+
+        Ok(Some((
+            recipient_id,
+            proto::Notification {
+                id: model.id.as_ref().to_proto(),
+                kind: proto.kind,
+                timestamp: model.created_at.as_ref().assume_utc().unix_timestamp() as u64,
+                is_read: false,
+                response: None,
+                content: proto.content,
+                entity_id: proto.entity_id,
+            },
+        )))
+    }
+
+    /// Remove an unread notification with the given recipient, kind and
+    /// entity id.
+    pub async fn remove_notification(
+        &self,
+        recipient_id: UserId,
+        notification: Notification,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<NotificationId>> {
+        let id = self
+            .find_notification(recipient_id, &notification, tx)
+            .await?;
+        if let Some(id) = id {
+            notification::Entity::delete_by_id(id).exec(tx).await?;
+        }
+        Ok(id)
+    }
+
+    /// Populate the response for the notification with the given kind and
+    /// entity id.
+    pub async fn mark_notification_as_read_with_response(
+        &self,
+        recipient_id: UserId,
+        notification: &Notification,
+        response: bool,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<(UserId, proto::Notification)>> {
+        self.mark_notification_as_read_internal(recipient_id, notification, Some(response), tx)
+            .await
+    }
+
+    pub async fn mark_notification_as_read(
+        &self,
+        recipient_id: UserId,
+        notification: &Notification,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<(UserId, proto::Notification)>> {
+        self.mark_notification_as_read_internal(recipient_id, notification, None, tx)
+            .await
+    }
+
+    pub async fn mark_notification_as_read_by_id(
+        &self,
+        recipient_id: UserId,
+        notification_id: NotificationId,
+    ) -> Result<NotificationBatch> {
+        self.transaction(|tx| async move {
+            let row = notification::Entity::update(notification::ActiveModel {
+                id: ActiveValue::Unchanged(notification_id),
+                recipient_id: ActiveValue::Unchanged(recipient_id),
+                is_read: ActiveValue::Set(true),
+                ..Default::default()
+            })
+            .exec(&*tx)
+            .await?;
+            Ok(model_to_proto(self, row)
+                .map(|notification| (recipient_id, notification))
+                .into_iter()
+                .collect())
+        })
+        .await
+    }
+
+    async fn mark_notification_as_read_internal(
+        &self,
+        recipient_id: UserId,
+        notification: &Notification,
+        response: Option<bool>,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<(UserId, proto::Notification)>> {
+        if let Some(id) = self
+            .find_notification(recipient_id, notification, &*tx)
+            .await?
+        {
+            let row = notification::Entity::update(notification::ActiveModel {
+                id: ActiveValue::Unchanged(id),
+                recipient_id: ActiveValue::Unchanged(recipient_id),
+                is_read: ActiveValue::Set(true),
+                response: if let Some(response) = response {
+                    ActiveValue::Set(Some(response))
+                } else {
+                    ActiveValue::NotSet
+                },
+                ..Default::default()
+            })
+            .exec(tx)
+            .await?;
+            Ok(model_to_proto(self, row).map(|notification| (recipient_id, notification)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Find an unread notification by its recipient, kind and entity id.
+    async fn find_notification(
+        &self,
+        recipient_id: UserId,
+        notification: &Notification,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<NotificationId>> {
+        let proto = notification.to_proto();
+        let kind = notification_kind_from_proto(self, &proto)?;
+
+        #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
+        enum QueryIds {
+            Id,
+        }
+
+        Ok(notification::Entity::find()
+            .select_only()
+            .column(notification::Column::Id)
+            .filter(
+                Condition::all()
+                    .add(notification::Column::RecipientId.eq(recipient_id))
+                    .add(notification::Column::IsRead.eq(false))
+                    .add(notification::Column::Kind.eq(kind))
+                    .add(if proto.entity_id.is_some() {
+                        notification::Column::EntityId.eq(proto.entity_id)
+                    } else {
+                        notification::Column::EntityId.is_null()
+                    }),
+            )
+            .into_values::<_, QueryIds>()
+            .one(&*tx)
+            .await?)
+    }
+}
+
+fn model_to_proto(this: &Database, row: notification::Model) -> Option<proto::Notification> {
+    let kind = this.notification_kinds_by_id.get(&row.kind)?;
+    Some(proto::Notification {
+        id: row.id.to_proto(),
+        kind: kind.to_string(),
+        timestamp: row.created_at.assume_utc().unix_timestamp() as u64,
+        is_read: row.is_read,
+        response: row.response,
+        content: row.content,
+        entity_id: row.entity_id.map(|id| id as u64),
+    })
+}
+
+fn notification_kind_from_proto(
+    this: &Database,
+    proto: &proto::Notification,
+) -> Result<NotificationKindId> {
+    Ok(this
+        .notification_kinds_by_name
+        .get(&proto.kind)
+        .copied()
+        .ok_or_else(|| anyhow!("invalid notification kind {:?}", proto.kind))?)
+}

crates/collab/src/db/queries/rooms.rs

@@ -50,10 +50,10 @@ impl Database {
                     .map(|participant| participant.user_id),
             );
 
-            let (channel_id, room) = self.get_channel_room(room_id, &tx).await?;
+            let (channel, room) = self.get_channel_room(room_id, &tx).await?;
             let channel_members;
-            if let Some(channel_id) = channel_id {
-                channel_members = self.get_channel_members_internal(channel_id, &tx).await?;
+            if let Some(channel) = &channel {
+                channel_members = self.get_channel_participants(channel, &tx).await?;
             } else {
                 channel_members = Vec::new();
 
@@ -69,7 +69,7 @@ impl Database {
 
             Ok(RefreshedRoom {
                 room,
-                channel_id,
+                channel_id: channel.map(|channel| channel.id),
                 channel_members,
                 stale_participant_user_ids,
                 canceled_calls_to_user_ids,
@@ -298,98 +298,137 @@ impl Database {
                 }
             }
 
-            #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
-            enum QueryParticipantIndices {
-                ParticipantIndex,
+            if channel_id.is_some() {
+                Err(anyhow!("tried to join channel call directly"))?
             }
-            let existing_participant_indices: Vec<i32> = room_participant::Entity::find()
-                .filter(
-                    room_participant::Column::RoomId
-                        .eq(room_id)
-                        .and(room_participant::Column::ParticipantIndex.is_not_null()),
-                )
-                .select_only()
-                .column(room_participant::Column::ParticipantIndex)
-                .into_values::<_, QueryParticipantIndices>()
-                .all(&*tx)
+
+            let participant_index = self
+                .get_next_participant_index_internal(room_id, &*tx)
                 .await?;
 
-            let mut participant_index = 0;
-            while existing_participant_indices.contains(&participant_index) {
-                participant_index += 1;
-            }
-
-            if let Some(channel_id) = channel_id {
-                self.check_user_is_channel_member(channel_id, user_id, &*tx)
-                    .await?;
-
-                room_participant::Entity::insert_many([room_participant::ActiveModel {
-                    room_id: ActiveValue::set(room_id),
-                    user_id: ActiveValue::set(user_id),
+            let result = room_participant::Entity::update_many()
+                .filter(
+                    Condition::all()
+                        .add(room_participant::Column::RoomId.eq(room_id))
+                        .add(room_participant::Column::UserId.eq(user_id))
+                        .add(room_participant::Column::AnsweringConnectionId.is_null()),
+                )
+                .set(room_participant::ActiveModel {
+                    participant_index: ActiveValue::Set(Some(participant_index)),
                     answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
                     answering_connection_server_id: ActiveValue::set(Some(ServerId(
                         connection.owner_id as i32,
                     ))),
                     answering_connection_lost: ActiveValue::set(false),
-                    calling_user_id: ActiveValue::set(user_id),
-                    calling_connection_id: ActiveValue::set(connection.id as i32),
-                    calling_connection_server_id: ActiveValue::set(Some(ServerId(
-                        connection.owner_id as i32,
-                    ))),
-                    participant_index: ActiveValue::Set(Some(participant_index)),
                     ..Default::default()
-                }])
-                .on_conflict(
-                    OnConflict::columns([room_participant::Column::UserId])
-                        .update_columns([
-                            room_participant::Column::AnsweringConnectionId,
-                            room_participant::Column::AnsweringConnectionServerId,
-                            room_participant::Column::AnsweringConnectionLost,
-                            room_participant::Column::ParticipantIndex,
-                        ])
-                        .to_owned(),
-                )
+                })
                 .exec(&*tx)
                 .await?;
-            } else {
-                let result = room_participant::Entity::update_many()
-                    .filter(
-                        Condition::all()
-                            .add(room_participant::Column::RoomId.eq(room_id))
-                            .add(room_participant::Column::UserId.eq(user_id))
-                            .add(room_participant::Column::AnsweringConnectionId.is_null()),
-                    )
-                    .set(room_participant::ActiveModel {
-                        participant_index: ActiveValue::Set(Some(participant_index)),
-                        answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
-                        answering_connection_server_id: ActiveValue::set(Some(ServerId(
-                            connection.owner_id as i32,
-                        ))),
-                        answering_connection_lost: ActiveValue::set(false),
-                        ..Default::default()
-                    })
-                    .exec(&*tx)
-                    .await?;
-                if result.rows_affected == 0 {
-                    Err(anyhow!("room does not exist or was already joined"))?;
-                }
+            if result.rows_affected == 0 {
+                Err(anyhow!("room does not exist or was already joined"))?;
             }
 
             let room = self.get_room(room_id, &tx).await?;
-            let channel_members = if let Some(channel_id) = channel_id {
-                self.get_channel_members_internal(channel_id, &tx).await?
-            } else {
-                Vec::new()
-            };
             Ok(JoinRoom {
                 room,
-                channel_id,
-                channel_members,
+                channel_id: None,
+                channel_members: vec![],
             })
         })
         .await
     }
 
+    async fn get_next_participant_index_internal(
+        &self,
+        room_id: RoomId,
+        tx: &DatabaseTransaction,
+    ) -> Result<i32> {
+        #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
+        enum QueryParticipantIndices {
+            ParticipantIndex,
+        }
+        let existing_participant_indices: Vec<i32> = room_participant::Entity::find()
+            .filter(
+                room_participant::Column::RoomId
+                    .eq(room_id)
+                    .and(room_participant::Column::ParticipantIndex.is_not_null()),
+            )
+            .select_only()
+            .column(room_participant::Column::ParticipantIndex)
+            .into_values::<_, QueryParticipantIndices>()
+            .all(&*tx)
+            .await?;
+
+        let mut participant_index = 0;
+        while existing_participant_indices.contains(&participant_index) {
+            participant_index += 1;
+        }
+
+        Ok(participant_index)
+    }
+
+    pub async fn channel_id_for_room(&self, room_id: RoomId) -> Result<Option<ChannelId>> {
+        self.transaction(|tx| async move {
+            let room: Option<room::Model> = room::Entity::find()
+                .filter(room::Column::Id.eq(room_id))
+                .one(&*tx)
+                .await?;
+
+            Ok(room.and_then(|room| room.channel_id))
+        })
+        .await
+    }
+
+    pub(crate) async fn join_channel_room_internal(
+        &self,
+        room_id: RoomId,
+        user_id: UserId,
+        connection: ConnectionId,
+        tx: &DatabaseTransaction,
+    ) -> Result<JoinRoom> {
+        let participant_index = self
+            .get_next_participant_index_internal(room_id, &*tx)
+            .await?;
+
+        room_participant::Entity::insert_many([room_participant::ActiveModel {
+            room_id: ActiveValue::set(room_id),
+            user_id: ActiveValue::set(user_id),
+            answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
+            answering_connection_server_id: ActiveValue::set(Some(ServerId(
+                connection.owner_id as i32,
+            ))),
+            answering_connection_lost: ActiveValue::set(false),
+            calling_user_id: ActiveValue::set(user_id),
+            calling_connection_id: ActiveValue::set(connection.id as i32),
+            calling_connection_server_id: ActiveValue::set(Some(ServerId(
+                connection.owner_id as i32,
+            ))),
+            participant_index: ActiveValue::Set(Some(participant_index)),
+            ..Default::default()
+        }])
+        .on_conflict(
+            OnConflict::columns([room_participant::Column::UserId])
+                .update_columns([
+                    room_participant::Column::AnsweringConnectionId,
+                    room_participant::Column::AnsweringConnectionServerId,
+                    room_participant::Column::AnsweringConnectionLost,
+                    room_participant::Column::ParticipantIndex,
+                ])
+                .to_owned(),
+        )
+        .exec(&*tx)
+        .await?;
+
+        let (channel, room) = self.get_channel_room(room_id, &tx).await?;
+        let channel = channel.ok_or_else(|| anyhow!("no channel for room"))?;
+        let channel_members = self.get_channel_participants(&channel, &*tx).await?;
+        Ok(JoinRoom {
+            room,
+            channel_id: Some(channel.id),
+            channel_members,
+        })
+    }
+
     pub async fn rejoin_room(
         &self,
         rejoin_room: proto::RejoinRoom,
@@ -679,16 +718,16 @@ impl Database {
                 });
             }
 
-            let (channel_id, room) = self.get_channel_room(room_id, &tx).await?;
-            let channel_members = if let Some(channel_id) = channel_id {
-                self.get_channel_members_internal(channel_id, &tx).await?
+            let (channel, room) = self.get_channel_room(room_id, &tx).await?;
+            let channel_members = if let Some(channel) = &channel {
+                self.get_channel_participants(&channel, &tx).await?
             } else {
                 Vec::new()
             };
 
             Ok(RejoinedRoom {
                 room,
-                channel_id,
+                channel_id: channel.map(|channel| channel.id),
                 channel_members,
                 rejoined_projects,
                 reshared_projects,
@@ -830,7 +869,7 @@ impl Database {
                     .exec(&*tx)
                     .await?;
 
-                let (channel_id, room) = self.get_channel_room(room_id, &tx).await?;
+                let (channel, room) = self.get_channel_room(room_id, &tx).await?;
                 let deleted = if room.participants.is_empty() {
                     let result = room::Entity::delete_by_id(room_id).exec(&*tx).await?;
                     result.rows_affected > 0
@@ -838,14 +877,14 @@ impl Database {
                     false
                 };
 
-                let channel_members = if let Some(channel_id) = channel_id {
-                    self.get_channel_members_internal(channel_id, &tx).await?
+                let channel_members = if let Some(channel) = &channel {
+                    self.get_channel_participants(channel, &tx).await?
                 } else {
                     Vec::new()
                 };
                 let left_room = LeftRoom {
                     room,
-                    channel_id,
+                    channel_id: channel.map(|channel| channel.id),
                     channel_members,
                     left_projects,
                     canceled_calls_to_user_ids,
@@ -1033,7 +1072,7 @@ impl Database {
         &self,
         room_id: RoomId,
         tx: &DatabaseTransaction,
-    ) -> Result<(Option<ChannelId>, proto::Room)> {
+    ) -> Result<(Option<channel::Model>, proto::Room)> {
         let db_room = room::Entity::find_by_id(room_id)
             .one(tx)
             .await?
@@ -1142,9 +1181,16 @@ impl Database {
                 project_id: db_follower.project_id.to_proto(),
             });
         }
+        drop(db_followers);
+
+        let channel = if let Some(channel_id) = db_room.channel_id {
+            Some(self.get_channel_internal(channel_id, &*tx).await?)
+        } else {
+            None
+        };
 
         Ok((
-            db_room.channel_id,
+            channel,
             proto::Room {
                 id: db_room.id.to_proto(),
                 live_kit_room: db_room.live_kit_room,

crates/collab/src/db/tables.rs

@@ -7,11 +7,13 @@ pub mod channel_buffer_collaborator;
 pub mod channel_chat_participant;
 pub mod channel_member;
 pub mod channel_message;
-pub mod channel_path;
+pub mod channel_message_mention;
 pub mod contact;
 pub mod feature_flag;
 pub mod follower;
 pub mod language_server;
+pub mod notification;
+pub mod notification_kind;
 pub mod observed_buffer_edits;
 pub mod observed_channel_messages;
 pub mod project;

crates/collab/src/db/tables/channel.rs

@@ -1,4 +1,4 @@
-use crate::db::ChannelId;
+use crate::db::{ChannelId, ChannelVisibility};
 use sea_orm::entity::prelude::*;
 
 #[derive(Clone, Debug, Default, PartialEq, Eq, DeriveEntityModel)]
@@ -7,6 +7,29 @@ pub struct Model {
     #[sea_orm(primary_key)]
     pub id: ChannelId,
     pub name: String,
+    pub visibility: ChannelVisibility,
+    pub parent_path: String,
+}
+
+impl Model {
+    pub fn parent_id(&self) -> Option<ChannelId> {
+        self.ancestors().last()
+    }
+
+    pub fn ancestors(&self) -> impl Iterator<Item = ChannelId> + '_ {
+        self.parent_path
+            .trim_end_matches('/')
+            .split('/')
+            .filter_map(|id| Some(ChannelId::from_proto(id.parse().ok()?)))
+    }
+
+    pub fn ancestors_including_self(&self) -> impl Iterator<Item = ChannelId> + '_ {
+        self.ancestors().chain(Some(self.id))
+    }
+
+    pub fn path(&self) -> String {
+        format!("{}{}/", self.parent_path, self.id)
+    }
 }
 
 impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/db/tables/channel_member.rs

@@ -1,7 +1,7 @@
-use crate::db::{channel_member, ChannelId, ChannelMemberId, UserId};
+use crate::db::{channel_member, ChannelId, ChannelMemberId, ChannelRole, UserId};
 use sea_orm::entity::prelude::*;
 
-#[derive(Clone, Debug, Default, PartialEq, Eq, DeriveEntityModel)]
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
 #[sea_orm(table_name = "channel_members")]
 pub struct Model {
     #[sea_orm(primary_key)]
@@ -9,7 +9,7 @@ pub struct Model {
     pub channel_id: ChannelId,
     pub user_id: UserId,
     pub accepted: bool,
-    pub admin: bool,
+    pub role: ChannelRole,
 }
 
 impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/db/tables/channel_message_mention.rs

@@ -0,0 +1,43 @@
+use crate::db::{MessageId, UserId};
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "channel_message_mentions")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub message_id: MessageId,
+    #[sea_orm(primary_key)]
+    pub start_offset: i32,
+    pub end_offset: i32,
+    pub user_id: UserId,
+}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::channel_message::Entity",
+        from = "Column::MessageId",
+        to = "super::channel_message::Column::Id"
+    )]
+    Message,
+    #[sea_orm(
+        belongs_to = "super::user::Entity",
+        from = "Column::UserId",
+        to = "super::user::Column::Id"
+    )]
+    MentionedUser,
+}
+
+impl Related<super::channel::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Message.def()
+    }
+}
+
+impl Related<super::user::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::MentionedUser.def()
+    }
+}

crates/collab/src/db/tables/notification.rs

@@ -0,0 +1,29 @@
+use crate::db::{NotificationId, NotificationKindId, UserId};
+use sea_orm::entity::prelude::*;
+use time::PrimitiveDateTime;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "notifications")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: NotificationId,
+    pub created_at: PrimitiveDateTime,
+    pub recipient_id: UserId,
+    pub kind: NotificationKindId,
+    pub entity_id: Option<i32>,
+    pub content: String,
+    pub is_read: bool,
+    pub response: Option<bool>,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::user::Entity",
+        from = "Column::RecipientId",
+        to = "super::user::Column::Id"
+    )]
+    Recipient,
+}
+
+impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/db/tables/notification_kind.rs

@@ -1,15 +1,15 @@
-use crate::db::ChannelId;
+use crate::db::NotificationKindId;
 use sea_orm::entity::prelude::*;
 
-#[derive(Clone, Debug, Default, PartialEq, Eq, DeriveEntityModel)]
-#[sea_orm(table_name = "channel_paths")]
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "notification_kinds")]
 pub struct Model {
     #[sea_orm(primary_key)]
-    pub id_path: String,
-    pub channel_id: ChannelId,
+    pub id: NotificationKindId,
+    pub name: String,
 }
 
-impl ActiveModelBehavior for ActiveModel {}
-
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
 pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/db/tests.rs

@@ -7,10 +7,12 @@ mod message_tests;
 use super::*;
 use gpui::executor::Background;
 use parking_lot::Mutex;
-use rpc::proto::ChannelEdge;
 use sea_orm::ConnectionTrait;
 use sqlx::migrate::MigrateDatabase;
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicI32, AtomicU32, Ordering::SeqCst},
+    Arc,
+};
 
 const TEST_RELEASE_CHANNEL: &'static str = "test";
 
@@ -31,7 +33,7 @@ impl TestDb {
         let mut db = runtime.block_on(async {
             let mut options = ConnectOptions::new(url);
             options.max_connections(5);
-            let db = Database::new(options, Executor::Deterministic(background))
+            let mut db = Database::new(options, Executor::Deterministic(background))
                 .await
                 .unwrap();
             let sql = include_str!(concat!(
@@ -45,6 +47,7 @@ impl TestDb {
                 ))
                 .await
                 .unwrap();
+            db.initialize_notification_kinds().await.unwrap();
             db
         });
 
@@ -79,11 +82,12 @@ impl TestDb {
             options
                 .max_connections(5)
                 .idle_timeout(Duration::from_secs(0));
-            let db = Database::new(options, Executor::Deterministic(background))
+            let mut db = Database::new(options, Executor::Deterministic(background))
                 .await
                 .unwrap();
             let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
             db.migrate(Path::new(migrations_path), false).await.unwrap();
+            db.initialize_notification_kinds().await.unwrap();
             db
         });
 
@@ -148,26 +152,39 @@ impl Drop for TestDb {
     }
 }
 
-/// The second tuples are (channel_id, parent)
-fn graph(channels: &[(ChannelId, &'static str)], edges: &[(ChannelId, ChannelId)]) -> ChannelGraph {
-    let mut graph = ChannelGraph {
-        channels: vec![],
-        edges: vec![],
-    };
-
-    for (id, name) in channels {
-        graph.channels.push(Channel {
+fn channel_tree(channels: &[(ChannelId, &[ChannelId], &'static str, ChannelRole)]) -> Vec<Channel> {
+    channels
+        .iter()
+        .map(|(id, parent_path, name, role)| Channel {
             id: *id,
             name: name.to_string(),
+            visibility: ChannelVisibility::Members,
+            role: *role,
+            parent_path: parent_path.to_vec(),
         })
-    }
+        .collect()
+}
+
+static GITHUB_USER_ID: AtomicI32 = AtomicI32::new(5);
+
+async fn new_test_user(db: &Arc<Database>, email: &str) -> UserId {
+    db.create_user(
+        email,
+        false,
+        NewUserParams {
+            github_login: email[0..email.find("@").unwrap()].to_string(),
+            github_user_id: GITHUB_USER_ID.fetch_add(1, SeqCst),
+        },
+    )
+    .await
+    .unwrap()
+    .user_id
+}
 
-    for (channel, parent) in edges {
-        graph.edges.push(ChannelEdge {
-            channel_id: channel.to_proto(),
-            parent_id: parent.to_proto(),
-        })
+static TEST_CONNECTION_ID: AtomicU32 = AtomicU32::new(1);
+fn new_test_connection(server: ServerId) -> ConnectionId {
+    ConnectionId {
+        id: TEST_CONNECTION_ID.fetch_add(1, SeqCst),
+        owner_id: server.0 as u32,
     }
-
-    graph
 }

crates/collab/src/db/tests/buffer_tests.rs

@@ -17,7 +17,6 @@ async fn test_channel_buffers(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "user_a".into(),
                 github_user_id: 101,
-                invite_count: 0,
             },
         )
         .await
@@ -30,7 +29,6 @@ async fn test_channel_buffers(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "user_b".into(),
                 github_user_id: 102,
-                invite_count: 0,
             },
         )
         .await
@@ -45,7 +43,6 @@ async fn test_channel_buffers(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "user_c".into(),
                 github_user_id: 102,
-                invite_count: 0,
             },
         )
         .await
@@ -56,7 +53,7 @@ async fn test_channel_buffers(db: &Arc<Database>) {
 
     let zed_id = db.create_root_channel("zed", a_id).await.unwrap();
 
-    db.invite_channel_member(zed_id, b_id, a_id, false)
+    db.invite_channel_member(zed_id, b_id, a_id, ChannelRole::Member)
         .await
         .unwrap();
 
@@ -178,7 +175,6 @@ async fn test_channel_buffers_last_operations(db: &Database) {
             NewUserParams {
                 github_login: "user_a".into(),
                 github_user_id: 101,
-                invite_count: 0,
             },
         )
         .await
@@ -191,7 +187,6 @@ async fn test_channel_buffers_last_operations(db: &Database) {
             NewUserParams {
                 github_login: "user_b".into(),
                 github_user_id: 102,
-                invite_count: 0,
             },
         )
         .await
@@ -211,7 +206,7 @@ async fn test_channel_buffers_last_operations(db: &Database) {
             .await
             .unwrap();
 
-        db.invite_channel_member(channel, observer_id, user_id, false)
+        db.invite_channel_member(channel, observer_id, user_id, ChannelRole::Member)
             .await
             .unwrap();
         db.respond_to_channel_invite(channel, observer_id, true)

crates/collab/src/db/tests/db_tests.rs

@@ -22,7 +22,6 @@ async fn test_get_users(db: &Arc<Database>) {
                 NewUserParams {
                     github_login: format!("user{i}"),
                     github_user_id: i,
-                    invite_count: 0,
                 },
             )
             .await
@@ -88,7 +87,6 @@ async fn test_get_or_create_user_by_github_account(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "login1".into(),
                 github_user_id: 101,
-                invite_count: 0,
             },
         )
         .await
@@ -101,7 +99,6 @@ async fn test_get_or_create_user_by_github_account(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "login2".into(),
                 github_user_id: 102,
-                invite_count: 0,
             },
         )
         .await
@@ -156,7 +153,6 @@ async fn test_create_access_tokens(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "u1".into(),
                 github_user_id: 1,
-                invite_count: 0,
             },
         )
         .await
@@ -238,7 +234,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
                 NewUserParams {
                     github_login: format!("user{i}"),
                     github_user_id: i,
-                    invite_count: 0,
                 },
             )
             .await
@@ -264,10 +259,7 @@ async fn test_add_contacts(db: &Arc<Database>) {
     );
     assert_eq!(
         db.get_contacts(user_2).await.unwrap(),
-        &[Contact::Incoming {
-            user_id: user_1,
-            should_notify: true
-        }]
+        &[Contact::Incoming { user_id: user_1 }]
     );
 
     // User 2 dismisses the contact request notification without accepting or rejecting.
@@ -280,10 +272,7 @@ async fn test_add_contacts(db: &Arc<Database>) {
         .unwrap();
     assert_eq!(
         db.get_contacts(user_2).await.unwrap(),
-        &[Contact::Incoming {
-            user_id: user_1,
-            should_notify: false
-        }]
+        &[Contact::Incoming { user_id: user_1 }]
     );
 
     // User can't accept their own contact request
@@ -299,7 +288,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_1).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_2,
-            should_notify: true,
             busy: false,
         }],
     );
@@ -309,7 +297,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_2).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_1,
-            should_notify: false,
             busy: false,
         }]
     );
@@ -326,7 +313,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_1).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_2,
-            should_notify: true,
             busy: false,
         }]
     );
@@ -339,7 +325,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_1).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_2,
-            should_notify: false,
             busy: false,
         }]
     );
@@ -353,12 +338,10 @@ async fn test_add_contacts(db: &Arc<Database>) {
         &[
             Contact::Accepted {
                 user_id: user_2,
-                should_notify: false,
                 busy: false,
             },
             Contact::Accepted {
                 user_id: user_3,
-                should_notify: false,
                 busy: false,
             }
         ]
@@ -367,7 +350,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_3).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_1,
-            should_notify: false,
             busy: false,
         }],
     );
@@ -383,7 +365,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_2).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_1,
-            should_notify: false,
             busy: false,
         }]
     );
@@ -391,7 +372,6 @@ async fn test_add_contacts(db: &Arc<Database>) {
         db.get_contacts(user_3).await.unwrap(),
         &[Contact::Accepted {
             user_id: user_1,
-            should_notify: false,
             busy: false,
         }],
     );
@@ -415,7 +395,6 @@ async fn test_metrics_id(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "person1".into(),
                 github_user_id: 101,
-                invite_count: 5,
             },
         )
         .await
@@ -431,7 +410,6 @@ async fn test_metrics_id(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "person2".into(),
                 github_user_id: 102,
-                invite_count: 5,
             },
         )
         .await
@@ -460,7 +438,6 @@ async fn test_project_count(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "admin".into(),
                 github_user_id: 0,
-                invite_count: 0,
             },
         )
         .await
@@ -472,7 +449,6 @@ async fn test_project_count(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "user".into(),
                 github_user_id: 1,
-                invite_count: 0,
             },
         )
         .await
@@ -554,7 +530,6 @@ async fn test_fuzzy_search_users() {
             NewUserParams {
                 github_login: github_login.into(),
                 github_user_id: i as i32,
-                invite_count: 0,
             },
         )
         .await
@@ -596,7 +571,6 @@ async fn test_non_matching_release_channels(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "admin".into(),
                 github_user_id: 0,
-                invite_count: 0,
             },
         )
         .await
@@ -608,7 +582,6 @@ async fn test_non_matching_release_channels(db: &Arc<Database>) {
             NewUserParams {
                 github_login: "user".into(),
                 github_user_id: 1,
-                invite_count: 0,
             },
         )
         .await

crates/collab/src/db/tests/feature_flag_tests.rs

@@ -18,7 +18,6 @@ async fn test_get_user_flags(db: &Arc<Database>) {
             NewUserParams {
                 github_login: format!("user1"),
                 github_user_id: 1,
-                invite_count: 0,
             },
         )
         .await
@@ -32,7 +31,6 @@ async fn test_get_user_flags(db: &Arc<Database>) {
             NewUserParams {
                 github_login: format!("user2"),
                 github_user_id: 2,
-                invite_count: 0,
             },
         )
         .await

crates/collab/src/db/tests/message_tests.rs

@@ -1,7 +1,9 @@
+use super::new_test_user;
 use crate::{
-    db::{Database, MessageId, NewUserParams},
+    db::{ChannelRole, Database, MessageId},
     test_both_dbs,
 };
+use channel::mentions_to_proto;
 use std::sync::Arc;
 use time::OffsetDateTime;
 
@@ -12,39 +14,38 @@ test_both_dbs!(
 );
 
 async fn test_channel_message_retrieval(db: &Arc<Database>) {
-    let user = db
-        .create_user(
-            "user@example.com",
-            false,
-            NewUserParams {
-                github_login: "user".into(),
-                github_user_id: 1,
-                invite_count: 0,
-            },
-        )
-        .await
-        .unwrap()
-        .user_id;
-    let channel = db.create_channel("channel", None, user).await.unwrap();
+    let user = new_test_user(db, "user@example.com").await;
+    let result = db.create_channel("channel", None, user).await.unwrap();
 
     let owner_id = db.create_server("test").await.unwrap().0 as u32;
-    db.join_channel_chat(channel, rpc::ConnectionId { owner_id, id: 0 }, user)
-        .await
-        .unwrap();
+    db.join_channel_chat(
+        result.channel.id,
+        rpc::ConnectionId { owner_id, id: 0 },
+        user,
+    )
+    .await
+    .unwrap();
 
     let mut all_messages = Vec::new();
     for i in 0..10 {
         all_messages.push(
-            db.create_channel_message(channel, user, &i.to_string(), OffsetDateTime::now_utc(), i)
-                .await
-                .unwrap()
-                .0
-                .to_proto(),
+            db.create_channel_message(
+                result.channel.id,
+                user,
+                &i.to_string(),
+                &[],
+                OffsetDateTime::now_utc(),
+                i,
+            )
+            .await
+            .unwrap()
+            .message_id
+            .to_proto(),
         );
     }
 
     let messages = db
-        .get_channel_messages(channel, user, 3, None)
+        .get_channel_messages(result.channel.id, user, 3, None)
         .await
         .unwrap()
         .into_iter()
@@ -54,7 +55,7 @@ async fn test_channel_message_retrieval(db: &Arc<Database>) {
 
     let messages = db
         .get_channel_messages(
-            channel,
+            result.channel.id,
             user,
             4,
             Some(MessageId::from_proto(all_messages[6])),
@@ -74,99 +75,154 @@ test_both_dbs!(
 );
 
 async fn test_channel_message_nonces(db: &Arc<Database>) {
-    let user = db
-        .create_user(
-            "user@example.com",
-            false,
-            NewUserParams {
-                github_login: "user".into(),
-                github_user_id: 1,
-                invite_count: 0,
-            },
+    let user_a = new_test_user(db, "user_a@example.com").await;
+    let user_b = new_test_user(db, "user_b@example.com").await;
+    let user_c = new_test_user(db, "user_c@example.com").await;
+    let channel = db.create_root_channel("channel", user_a).await.unwrap();
+    db.invite_channel_member(channel, user_b, user_a, ChannelRole::Member)
+        .await
+        .unwrap();
+    db.invite_channel_member(channel, user_c, user_a, ChannelRole::Member)
+        .await
+        .unwrap();
+    db.respond_to_channel_invite(channel, user_b, true)
+        .await
+        .unwrap();
+    db.respond_to_channel_invite(channel, user_c, true)
+        .await
+        .unwrap();
+
+    let owner_id = db.create_server("test").await.unwrap().0 as u32;
+    db.join_channel_chat(channel, rpc::ConnectionId { owner_id, id: 0 }, user_a)
+        .await
+        .unwrap();
+    db.join_channel_chat(channel, rpc::ConnectionId { owner_id, id: 1 }, user_b)
+        .await
+        .unwrap();
+
+    // As user A, create messages that re-use the same nonces. The requests
+    // succeed, but return the same ids.
+    let id1 = db
+        .create_channel_message(
+            channel,
+            user_a,
+            "hi @user_b",
+            &mentions_to_proto(&[(3..10, user_b.to_proto())]),
+            OffsetDateTime::now_utc(),
+            100,
         )
         .await
         .unwrap()
-        .user_id;
-    let channel = db.create_channel("channel", None, user).await.unwrap();
+        .message_id;
+    let id2 = db
+        .create_channel_message(
+            channel,
+            user_a,
+            "hello, fellow users",
+            &mentions_to_proto(&[]),
+            OffsetDateTime::now_utc(),
+            200,
+        )
+        .await
+        .unwrap()
+        .message_id;
+    let id3 = db
+        .create_channel_message(
+            channel,
+            user_a,
+            "bye @user_c (same nonce as first message)",
+            &mentions_to_proto(&[(4..11, user_c.to_proto())]),
+            OffsetDateTime::now_utc(),
+            100,
+        )
+        .await
+        .unwrap()
+        .message_id;
+    let id4 = db
+        .create_channel_message(
+            channel,
+            user_a,
+            "omg (same nonce as second message)",
+            &mentions_to_proto(&[]),
+            OffsetDateTime::now_utc(),
+            200,
+        )
+        .await
+        .unwrap()
+        .message_id;
 
-    let owner_id = db.create_server("test").await.unwrap().0 as u32;
+    // As a different user, reuse one of the same nonces. This request succeeds
+    // and returns a different id.
+    let id5 = db
+        .create_channel_message(
+            channel,
+            user_b,
+            "omg @user_a (same nonce as user_a's first message)",
+            &mentions_to_proto(&[(4..11, user_a.to_proto())]),
+            OffsetDateTime::now_utc(),
+            100,
+        )
+        .await
+        .unwrap()
+        .message_id;
 
-    db.join_channel_chat(channel, rpc::ConnectionId { owner_id, id: 0 }, user)
-        .await
-        .unwrap();
+    assert_ne!(id1, id2);
+    assert_eq!(id1, id3);
+    assert_eq!(id2, id4);
+    assert_ne!(id5, id1);
 
-    let msg1_id = db
-        .create_channel_message(channel, user, "1", OffsetDateTime::now_utc(), 1)
+    let messages = db
+        .get_channel_messages(channel, user_a, 5, None)
         .await
-        .unwrap();
-    let msg2_id = db
-        .create_channel_message(channel, user, "2", OffsetDateTime::now_utc(), 2)
-        .await
-        .unwrap();
-    let msg3_id = db
-        .create_channel_message(channel, user, "3", OffsetDateTime::now_utc(), 1)
-        .await
-        .unwrap();
-    let msg4_id = db
-        .create_channel_message(channel, user, "4", OffsetDateTime::now_utc(), 2)
-        .await
-        .unwrap();
-
-    assert_ne!(msg1_id, msg2_id);
-    assert_eq!(msg1_id, msg3_id);
-    assert_eq!(msg2_id, msg4_id);
+        .unwrap()
+        .into_iter()
+        .map(|m| (m.id, m.body, m.mentions))
+        .collect::<Vec<_>>();
+    assert_eq!(
+        messages,
+        &[
+            (
+                id1.to_proto(),
+                "hi @user_b".into(),
+                mentions_to_proto(&[(3..10, user_b.to_proto())]),
+            ),
+            (
+                id2.to_proto(),
+                "hello, fellow users".into(),
+                mentions_to_proto(&[])
+            ),
+            (
+                id5.to_proto(),
+                "omg @user_a (same nonce as user_a's first message)".into(),
+                mentions_to_proto(&[(4..11, user_a.to_proto())]),
+            ),
+        ]
+    );
 }
 
 test_both_dbs!(
-    test_channel_message_new_notification,
-    test_channel_message_new_notification_postgres,
-    test_channel_message_new_notification_sqlite
+    test_unseen_channel_messages,
+    test_unseen_channel_messages_postgres,
+    test_unseen_channel_messages_sqlite
 );
 
-async fn test_channel_message_new_notification(db: &Arc<Database>) {
-    let user = db
-        .create_user(
-            "user_a@example.com",
-            false,
-            NewUserParams {
-                github_login: "user_a".into(),
-                github_user_id: 1,
-                invite_count: 0,
-            },
-        )
+async fn test_unseen_channel_messages(db: &Arc<Database>) {
+    let user = new_test_user(db, "user_a@example.com").await;
+    let observer = new_test_user(db, "user_b@example.com").await;
+
+    let channel_1 = db.create_root_channel("channel", user).await.unwrap();
+    let channel_2 = db.create_root_channel("channel-2", user).await.unwrap();
+
+    db.invite_channel_member(channel_1, observer, user, ChannelRole::Member)
         .await
-        .unwrap()
-        .user_id;
-    let observer = db
-        .create_user(
-            "user_b@example.com",
-            false,
-            NewUserParams {
-                github_login: "user_b".into(),
-                github_user_id: 1,
-                invite_count: 0,
-            },
-        )
-        .await
-        .unwrap()
-        .user_id;
-
-    let channel_1 = db.create_channel("channel", None, user).await.unwrap();
-
-    let channel_2 = db.create_channel("channel-2", None, user).await.unwrap();
-
-    db.invite_channel_member(channel_1, observer, user, false)
+        .unwrap();
+    db.invite_channel_member(channel_2, observer, user, ChannelRole::Member)
         .await
         .unwrap();
 
     db.respond_to_channel_invite(channel_1, observer, true)
         .await
         .unwrap();
-
-    db.invite_channel_member(channel_2, observer, user, false)
-        .await
-        .unwrap();
-
     db.respond_to_channel_invite(channel_2, observer, true)
         .await
         .unwrap();
@@ -179,28 +235,31 @@ async fn test_channel_message_new_notification(db: &Arc<Database>) {
         .unwrap();
 
     let _ = db
-        .create_channel_message(channel_1, user, "1_1", OffsetDateTime::now_utc(), 1)
+        .create_channel_message(channel_1, user, "1_1", &[], OffsetDateTime::now_utc(), 1)
         .await
         .unwrap();
 
-    let (second_message, _, _) = db
-        .create_channel_message(channel_1, user, "1_2", OffsetDateTime::now_utc(), 2)
+    let second_message = db
+        .create_channel_message(channel_1, user, "1_2", &[], OffsetDateTime::now_utc(), 2)
         .await
-        .unwrap();
+        .unwrap()
+        .message_id;
 
-    let (third_message, _, _) = db
-        .create_channel_message(channel_1, user, "1_3", OffsetDateTime::now_utc(), 3)
+    let third_message = db
+        .create_channel_message(channel_1, user, "1_3", &[], OffsetDateTime::now_utc(), 3)
         .await
-        .unwrap();
+        .unwrap()
+        .message_id;
 
     db.join_channel_chat(channel_2, user_connection_id, user)
         .await
         .unwrap();
 
-    let (fourth_message, _, _) = db
-        .create_channel_message(channel_2, user, "2_1", OffsetDateTime::now_utc(), 4)
+    let fourth_message = db
+        .create_channel_message(channel_2, user, "2_1", &[], OffsetDateTime::now_utc(), 4)
         .await
-        .unwrap();
+        .unwrap()
+        .message_id;
 
     // Check that observer has new messages
     let unseen_messages = db
@@ -295,3 +354,101 @@ async fn test_channel_message_new_notification(db: &Arc<Database>) {
         }]
     );
 }
+
+test_both_dbs!(
+    test_channel_message_mentions,
+    test_channel_message_mentions_postgres,
+    test_channel_message_mentions_sqlite
+);
+
+async fn test_channel_message_mentions(db: &Arc<Database>) {
+    let user_a = new_test_user(db, "user_a@example.com").await;
+    let user_b = new_test_user(db, "user_b@example.com").await;
+    let user_c = new_test_user(db, "user_c@example.com").await;
+
+    let channel = db
+        .create_channel("channel", None, user_a)
+        .await
+        .unwrap()
+        .channel
+        .id;
+    db.invite_channel_member(channel, user_b, user_a, ChannelRole::Member)
+        .await
+        .unwrap();
+    db.respond_to_channel_invite(channel, user_b, true)
+        .await
+        .unwrap();
+
+    let owner_id = db.create_server("test").await.unwrap().0 as u32;
+    let connection_id = rpc::ConnectionId { owner_id, id: 0 };
+    db.join_channel_chat(channel, connection_id, user_a)
+        .await
+        .unwrap();
+
+    db.create_channel_message(
+        channel,
+        user_a,
+        "hi @user_b and @user_c",
+        &mentions_to_proto(&[(3..10, user_b.to_proto()), (15..22, user_c.to_proto())]),
+        OffsetDateTime::now_utc(),
+        1,
+    )
+    .await
+    .unwrap();
+    db.create_channel_message(
+        channel,
+        user_a,
+        "bye @user_c",
+        &mentions_to_proto(&[(4..11, user_c.to_proto())]),
+        OffsetDateTime::now_utc(),
+        2,
+    )
+    .await
+    .unwrap();
+    db.create_channel_message(
+        channel,
+        user_a,
+        "umm",
+        &mentions_to_proto(&[]),
+        OffsetDateTime::now_utc(),
+        3,
+    )
+    .await
+    .unwrap();
+    db.create_channel_message(
+        channel,
+        user_a,
+        "@user_b, stop.",
+        &mentions_to_proto(&[(0..7, user_b.to_proto())]),
+        OffsetDateTime::now_utc(),
+        4,
+    )
+    .await
+    .unwrap();
+
+    let messages = db
+        .get_channel_messages(channel, user_b, 5, None)
+        .await
+        .unwrap()
+        .into_iter()
+        .map(|m| (m.body, m.mentions))
+        .collect::<Vec<_>>();
+    assert_eq!(
+        &messages,
+        &[
+            (
+                "hi @user_b and @user_c".into(),
+                mentions_to_proto(&[(3..10, user_b.to_proto()), (15..22, user_c.to_proto())]),
+            ),
+            (
+                "bye @user_c".into(),
+                mentions_to_proto(&[(4..11, user_c.to_proto())]),
+            ),
+            ("umm".into(), mentions_to_proto(&[]),),
+            (
+                "@user_b, stop.".into(),
+                mentions_to_proto(&[(0..7, user_b.to_proto())]),
+            ),
+        ]
+    );
+}

crates/collab/src/lib.rs

@@ -119,7 +119,9 @@ impl AppState {
     pub async fn new(config: Config) -> Result<Arc<Self>> {
         let mut db_options = db::ConnectOptions::new(config.database_url.clone());
         db_options.max_connections(config.database_max_connections);
-        let db = Database::new(db_options, Executor::Production).await?;
+        let mut db = Database::new(db_options, Executor::Production).await?;
+        db.initialize_notification_kinds().await?;
+
         let live_kit_client = if let Some(((server, key), secret)) = config
             .live_kit_server
             .as_ref()

crates/collab/src/tests.rs

@@ -6,6 +6,7 @@ mod channel_message_tests;
 mod channel_tests;
 mod following_tests;
 mod integration_tests;
+mod notification_tests;
 mod random_channel_buffer_tests;
 mod random_project_collaboration_tests;
 mod randomized_test_helpers;
@@ -39,3 +40,7 @@ fn room_participants(room: &ModelHandle<Room>, cx: &mut TestAppContext) -> RoomP
         RoomParticipants { remote, pending }
     })
 }
+
+fn channel_id(room: &ModelHandle<Room>, cx: &mut TestAppContext) -> Option<u64> {
+    cx.read(|cx| room.read(cx).channel_id())
+}

crates/collab/src/tests/channel_buffer_tests.rs

@@ -3,7 +3,7 @@ use crate::{
     tests::TestServer,
 };
 use call::ActiveCall;
-use channel::{Channel, ACKNOWLEDGE_DEBOUNCE_INTERVAL};
+use channel::ACKNOWLEDGE_DEBOUNCE_INTERVAL;
 use client::ParticipantIndex;
 use client::{Collaborator, UserId};
 use collab_ui::channel_view::ChannelView;
@@ -407,11 +407,8 @@ async fn test_channel_buffer_disconnect(
     server.disconnect_client(client_a.peer_id().unwrap());
     deterministic.advance_clock(RECEIVE_TIMEOUT + RECONNECT_TIMEOUT);
 
-    channel_buffer_a.update(cx_a, |buffer, _| {
-        assert_eq!(
-            buffer.channel().as_ref(),
-            &channel(channel_id, "the-channel")
-        );
+    channel_buffer_a.update(cx_a, |buffer, cx| {
+        assert_eq!(buffer.channel(cx).unwrap().name, "the-channel");
         assert!(!buffer.is_connected());
     });
 
@@ -432,24 +429,12 @@ async fn test_channel_buffer_disconnect(
     deterministic.run_until_parked();
 
     // Channel buffer observed the deletion
-    channel_buffer_b.update(cx_b, |buffer, _| {
-        assert_eq!(
-            buffer.channel().as_ref(),
-            &channel(channel_id, "the-channel")
-        );
+    channel_buffer_b.update(cx_b, |buffer, cx| {
+        assert!(buffer.channel(cx).is_none());
         assert!(!buffer.is_connected());
     });
 }
 
-fn channel(id: u64, name: &'static str) -> Channel {
-    Channel {
-        id,
-        name: name.to_string(),
-        unseen_note_version: None,
-        unseen_message_id: None,
-    }
-}
-
 #[gpui::test]
 async fn test_rejoin_channel_buffer(
     deterministic: Arc<Deterministic>,
@@ -694,7 +679,7 @@ async fn test_following_to_channel_notes_without_a_shared_project(
         .await
         .unwrap();
     channel_view_1_a.update(cx_a, |notes, cx| {
-        assert_eq!(notes.channel(cx).name, "channel-1");
+        assert_eq!(notes.channel(cx).unwrap().name, "channel-1");
         notes.editor.update(cx, |editor, cx| {
             editor.insert("Hello from A.", cx);
             editor.change_selections(None, cx, |selections| {
@@ -726,7 +711,7 @@ async fn test_following_to_channel_notes_without_a_shared_project(
             .expect("active item is not a channel view")
     });
     channel_view_1_b.read_with(cx_b, |notes, cx| {
-        assert_eq!(notes.channel(cx).name, "channel-1");
+        assert_eq!(notes.channel(cx).unwrap().name, "channel-1");
         let editor = notes.editor.read(cx);
         assert_eq!(editor.text(cx), "Hello from A.");
         assert_eq!(editor.selections.ranges::<usize>(cx), &[3..4]);
@@ -738,7 +723,7 @@ async fn test_following_to_channel_notes_without_a_shared_project(
         .await
         .unwrap();
     channel_view_2_a.read_with(cx_a, |notes, cx| {
-        assert_eq!(notes.channel(cx).name, "channel-2");
+        assert_eq!(notes.channel(cx).unwrap().name, "channel-2");
     });
 
     // Client B is taken to the notes for channel 2.
@@ -755,7 +740,7 @@ async fn test_following_to_channel_notes_without_a_shared_project(
             .expect("active item is not a channel view")
     });
     channel_view_2_b.read_with(cx_b, |notes, cx| {
-        assert_eq!(notes.channel(cx).name, "channel-2");
+        assert_eq!(notes.channel(cx).unwrap().name, "channel-2");
     });
 }
 

crates/collab/src/tests/channel_message_tests.rs

@@ -1,27 +1,30 @@
 use crate::{rpc::RECONNECT_TIMEOUT, tests::TestServer};
-use channel::{ChannelChat, ChannelMessageId};
+use channel::{ChannelChat, ChannelMessageId, MessageParams};
 use collab_ui::chat_panel::ChatPanel;
 use gpui::{executor::Deterministic, BorrowAppContext, ModelHandle, TestAppContext};
+use rpc::Notification;
 use std::sync::Arc;
 use workspace::dock::Panel;
 
 #[gpui::test]
 async fn test_basic_channel_messages(
     deterministic: Arc<Deterministic>,
-    cx_a: &mut TestAppContext,
-    cx_b: &mut TestAppContext,
+    mut cx_a: &mut TestAppContext,
+    mut cx_b: &mut TestAppContext,
+    mut cx_c: &mut TestAppContext,
 ) {
     deterministic.forbid_parking();
     let mut server = TestServer::start(&deterministic).await;
     let client_a = server.create_client(cx_a, "user_a").await;
     let client_b = server.create_client(cx_b, "user_b").await;
+    let client_c = server.create_client(cx_c, "user_c").await;
 
     let channel_id = server
         .make_channel(
             "the-channel",
             None,
             (&client_a, cx_a),
-            &mut [(&client_b, cx_b)],
+            &mut [(&client_b, cx_b), (&client_c, cx_c)],
         )
         .await;
 
@@ -36,8 +39,17 @@ async fn test_basic_channel_messages(
         .await
         .unwrap();
 
-    channel_chat_a
-        .update(cx_a, |c, cx| c.send_message("one".into(), cx).unwrap())
+    let message_id = channel_chat_a
+        .update(cx_a, |c, cx| {
+            c.send_message(
+                MessageParams {
+                    text: "hi @user_c!".into(),
+                    mentions: vec![(3..10, client_c.id())],
+                },
+                cx,
+            )
+            .unwrap()
+        })
         .await
         .unwrap();
     channel_chat_a
@@ -52,15 +64,55 @@ async fn test_basic_channel_messages(
         .unwrap();
 
     deterministic.run_until_parked();
-    channel_chat_a.update(cx_a, |c, _| {
+
+    let channel_chat_c = client_c
+        .channel_store()
+        .update(cx_c, |store, cx| store.open_channel_chat(channel_id, cx))
+        .await
+        .unwrap();
+
+    for (chat, cx) in [
+        (&channel_chat_a, &mut cx_a),
+        (&channel_chat_b, &mut cx_b),
+        (&channel_chat_c, &mut cx_c),
+    ] {
+        chat.update(*cx, |c, _| {
+            assert_eq!(
+                c.messages()
+                    .iter()
+                    .map(|m| (m.body.as_str(), m.mentions.as_slice()))
+                    .collect::<Vec<_>>(),
+                vec![
+                    ("hi @user_c!", [(3..10, client_c.id())].as_slice()),
+                    ("two", &[]),
+                    ("three", &[])
+                ],
+                "results for user {}",
+                c.client().id(),
+            );
+        });
+    }
+
+    client_c.notification_store().update(cx_c, |store, _| {
+        assert_eq!(store.notification_count(), 2);
+        assert_eq!(store.unread_notification_count(), 1);
         assert_eq!(
-            c.messages()
-                .iter()
-                .map(|m| m.body.as_str())
-                .collect::<Vec<_>>(),
-            vec!["one", "two", "three"]
+            store.notification_at(0).unwrap().notification,
+            Notification::ChannelMessageMention {
+                message_id,
+                sender_id: client_a.id(),
+                channel_id,
+            }
         );
-    })
+        assert_eq!(
+            store.notification_at(1).unwrap().notification,
+            Notification::ChannelInvitation {
+                channel_id,
+                channel_name: "the-channel".to_string(),
+                inviter_id: client_a.id()
+            }
+        );
+    });
 }
 
 #[gpui::test]
@@ -280,7 +332,7 @@ async fn test_channel_message_changes(
     chat_panel_b
         .update(cx_b, |chat_panel, cx| {
             chat_panel.set_active(true, cx);
-            chat_panel.select_channel(channel_id, cx)
+            chat_panel.select_channel(channel_id, None, cx)
         })
         .await
         .unwrap();

crates/collab/src/tests/following_tests.rs

@@ -1,6 +1,6 @@
 use crate::{rpc::RECONNECT_TIMEOUT, tests::TestServer};
 use call::ActiveCall;
-use collab_ui::project_shared_notification::ProjectSharedNotification;
+use collab_ui::notifications::project_shared_notification::ProjectSharedNotification;
 use editor::{Editor, ExcerptRange, MultiBuffer};
 use gpui::{executor::Deterministic, geometry::vector::vec2f, TestAppContext, ViewHandle};
 use live_kit_client::MacOSDisplay;

crates/collab/src/tests/integration_tests.rs

@@ -1,6 +1,6 @@
 use crate::{
     rpc::{CLEANUP_TIMEOUT, RECONNECT_TIMEOUT},
-    tests::{room_participants, RoomParticipants, TestClient, TestServer},
+    tests::{channel_id, room_participants, RoomParticipants, TestClient, TestServer},
 };
 use call::{room, ActiveCall, ParticipantLocation, Room};
 use client::{User, RECEIVE_TIMEOUT};
@@ -15,8 +15,8 @@ use gpui::{executor::Deterministic, test::EmptyView, AppContext, ModelHandle, Te
 use indoc::indoc;
 use language::{
     language_settings::{AllLanguageSettings, Formatter, InlayHintSettings},
-    tree_sitter_rust, Anchor, BundledFormatter, Diagnostic, DiagnosticEntry, FakeLspAdapter,
-    Language, LanguageConfig, LineEnding, OffsetRangeExt, Point, Rope,
+    tree_sitter_rust, Anchor, Diagnostic, DiagnosticEntry, FakeLspAdapter, Language,
+    LanguageConfig, LineEnding, OffsetRangeExt, Point, Rope,
 };
 use live_kit_client::MacOSDisplay;
 use lsp::LanguageServerId;
@@ -469,6 +469,119 @@ async fn test_calling_multiple_users_simultaneously(
     );
 }
 
+#[gpui::test(iterations = 10)]
+async fn test_joining_channels_and_calling_multiple_users_simultaneously(
+    deterministic: Arc<Deterministic>,
+    cx_a: &mut TestAppContext,
+    cx_b: &mut TestAppContext,
+    cx_c: &mut TestAppContext,
+) {
+    deterministic.forbid_parking();
+    let mut server = TestServer::start(&deterministic).await;
+
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let client_b = server.create_client(cx_b, "user_b").await;
+    let client_c = server.create_client(cx_c, "user_c").await;
+    server
+        .make_contacts(&mut [(&client_a, cx_a), (&client_b, cx_b), (&client_c, cx_c)])
+        .await;
+
+    let channel_1 = server
+        .make_channel(
+            "channel1",
+            None,
+            (&client_a, cx_a),
+            &mut [(&client_b, cx_b), (&client_c, cx_c)],
+        )
+        .await;
+
+    let channel_2 = server
+        .make_channel(
+            "channel2",
+            None,
+            (&client_a, cx_a),
+            &mut [(&client_b, cx_b), (&client_c, cx_c)],
+        )
+        .await;
+
+    let active_call_a = cx_a.read(ActiveCall::global);
+
+    // Simultaneously join channel 1 and then channel 2
+    active_call_a
+        .update(cx_a, |call, cx| call.join_channel(channel_1, cx))
+        .detach();
+    let join_channel_2 = active_call_a.update(cx_a, |call, cx| call.join_channel(channel_2, cx));
+
+    join_channel_2.await.unwrap();
+
+    let room_a = active_call_a.read_with(cx_a, |call, _| call.room().unwrap().clone());
+    deterministic.run_until_parked();
+
+    assert_eq!(channel_id(&room_a, cx_a), Some(channel_2));
+
+    // Leave the room
+    active_call_a
+        .update(cx_a, |call, cx| {
+            let hang_up = call.hang_up(cx);
+            hang_up
+        })
+        .await
+        .unwrap();
+
+    // Initiating invites and then joining a channel should fail gracefully
+    let b_invite = active_call_a.update(cx_a, |call, cx| {
+        call.invite(client_b.user_id().unwrap(), None, cx)
+    });
+    let c_invite = active_call_a.update(cx_a, |call, cx| {
+        call.invite(client_c.user_id().unwrap(), None, cx)
+    });
+
+    let join_channel = active_call_a.update(cx_a, |call, cx| call.join_channel(channel_1, cx));
+
+    b_invite.await.unwrap();
+    c_invite.await.unwrap();
+    join_channel.await.unwrap();
+
+    let room_a = active_call_a.read_with(cx_a, |call, _| call.room().unwrap().clone());
+    deterministic.run_until_parked();
+
+    assert_eq!(
+        room_participants(&room_a, cx_a),
+        RoomParticipants {
+            remote: Default::default(),
+            pending: vec!["user_b".to_string(), "user_c".to_string()]
+        }
+    );
+
+    assert_eq!(channel_id(&room_a, cx_a), None);
+
+    // Leave the room
+    active_call_a
+        .update(cx_a, |call, cx| {
+            let hang_up = call.hang_up(cx);
+            hang_up
+        })
+        .await
+        .unwrap();
+
+    // Simultaneously join channel 1 and call user B and user C from client A.
+    let join_channel = active_call_a.update(cx_a, |call, cx| call.join_channel(channel_1, cx));
+
+    let b_invite = active_call_a.update(cx_a, |call, cx| {
+        call.invite(client_b.user_id().unwrap(), None, cx)
+    });
+    let c_invite = active_call_a.update(cx_a, |call, cx| {
+        call.invite(client_c.user_id().unwrap(), None, cx)
+    });
+
+    join_channel.await.unwrap();
+    b_invite.await.unwrap();
+    c_invite.await.unwrap();
+
+    active_call_a.read_with(cx_a, |call, _| call.room().unwrap().clone());
+    deterministic.run_until_parked();
+}
+
 #[gpui::test(iterations = 10)]
 async fn test_room_uniqueness(
     deterministic: Arc<Deterministic>,
@@ -4530,6 +4643,7 @@ async fn test_prettier_formatting_buffer(
         LanguageConfig {
             name: "Rust".into(),
             path_suffixes: vec!["rs".to_string()],
+            prettier_parser_name: Some("test_parser".to_string()),
             ..Default::default()
         },
         Some(tree_sitter_rust::language()),
@@ -4537,10 +4651,7 @@ async fn test_prettier_formatting_buffer(
     let test_plugin = "test_plugin";
     let mut fake_language_servers = language
         .set_fake_lsp_adapter(Arc::new(FakeLspAdapter {
-            enabled_formatters: vec![BundledFormatter::Prettier {
-                parser_name: Some("test_parser"),
-                plugin_names: vec![test_plugin],
-            }],
+            prettier_plugins: vec![test_plugin],
             ..Default::default()
         }))
         .await;
@@ -4557,11 +4668,7 @@ async fn test_prettier_formatting_buffer(
         .insert_tree(&directory, json!({ "a.rs": buffer_text }))
         .await;
     let (project_a, worktree_id) = client_a.build_local_project(&directory, cx_a).await;
-    let prettier_format_suffix = project_a.update(cx_a, |project, _| {
-        let suffix = project.enable_test_prettier(&[test_plugin]);
-        project.languages().add(language);
-        suffix
-    });
+    let prettier_format_suffix = project::TEST_PRETTIER_FORMAT_SUFFIX;
     let buffer_a = cx_a
         .background()
         .spawn(project_a.update(cx_a, |p, cx| p.open_buffer((worktree_id, "a.rs"), cx)))

crates/collab/src/tests/notification_tests.rs

@@ -0,0 +1,159 @@
+use crate::tests::TestServer;
+use gpui::{executor::Deterministic, TestAppContext};
+use notifications::NotificationEvent;
+use parking_lot::Mutex;
+use rpc::{proto, Notification};
+use std::sync::Arc;
+
+#[gpui::test]
+async fn test_notifications(
+    deterministic: Arc<Deterministic>,
+    cx_a: &mut TestAppContext,
+    cx_b: &mut TestAppContext,
+) {
+    deterministic.forbid_parking();
+    let mut server = TestServer::start(&deterministic).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let client_b = server.create_client(cx_b, "user_b").await;
+
+    let notification_events_a = Arc::new(Mutex::new(Vec::new()));
+    let notification_events_b = Arc::new(Mutex::new(Vec::new()));
+    client_a.notification_store().update(cx_a, |_, cx| {
+        let events = notification_events_a.clone();
+        cx.subscribe(&cx.handle(), move |_, _, event, _| {
+            events.lock().push(event.clone());
+        })
+        .detach()
+    });
+    client_b.notification_store().update(cx_b, |_, cx| {
+        let events = notification_events_b.clone();
+        cx.subscribe(&cx.handle(), move |_, _, event, _| {
+            events.lock().push(event.clone());
+        })
+        .detach()
+    });
+
+    // Client A sends a contact request to client B.
+    client_a
+        .user_store()
+        .update(cx_a, |store, cx| store.request_contact(client_b.id(), cx))
+        .await
+        .unwrap();
+
+    // Client B receives a contact request notification and responds to the
+    // request, accepting it.
+    deterministic.run_until_parked();
+    client_b.notification_store().update(cx_b, |store, cx| {
+        assert_eq!(store.notification_count(), 1);
+        assert_eq!(store.unread_notification_count(), 1);
+
+        let entry = store.notification_at(0).unwrap();
+        assert_eq!(
+            entry.notification,
+            Notification::ContactRequest {
+                sender_id: client_a.id()
+            }
+        );
+        assert!(!entry.is_read);
+        assert_eq!(
+            &notification_events_b.lock()[0..],
+            &[
+                NotificationEvent::NewNotification {
+                    entry: entry.clone(),
+                },
+                NotificationEvent::NotificationsUpdated {
+                    old_range: 0..0,
+                    new_count: 1
+                }
+            ]
+        );
+
+        store.respond_to_notification(entry.notification.clone(), true, cx);
+    });
+
+    // Client B sees the notification is now read, and that they responded.
+    deterministic.run_until_parked();
+    client_b.notification_store().read_with(cx_b, |store, _| {
+        assert_eq!(store.notification_count(), 1);
+        assert_eq!(store.unread_notification_count(), 0);
+
+        let entry = store.notification_at(0).unwrap();
+        assert!(entry.is_read);
+        assert_eq!(entry.response, Some(true));
+        assert_eq!(
+            &notification_events_b.lock()[2..],
+            &[
+                NotificationEvent::NotificationRead {
+                    entry: entry.clone(),
+                },
+                NotificationEvent::NotificationsUpdated {
+                    old_range: 0..1,
+                    new_count: 1
+                }
+            ]
+        );
+    });
+
+    // Client A receives a notification that client B accepted their request.
+    client_a.notification_store().read_with(cx_a, |store, _| {
+        assert_eq!(store.notification_count(), 1);
+        assert_eq!(store.unread_notification_count(), 1);
+
+        let entry = store.notification_at(0).unwrap();
+        assert_eq!(
+            entry.notification,
+            Notification::ContactRequestAccepted {
+                responder_id: client_b.id()
+            }
+        );
+        assert!(!entry.is_read);
+    });
+
+    // Client A creates a channel and invites client B to be a member.
+    let channel_id = client_a
+        .channel_store()
+        .update(cx_a, |store, cx| {
+            store.create_channel("the-channel", None, cx)
+        })
+        .await
+        .unwrap();
+    client_a
+        .channel_store()
+        .update(cx_a, |store, cx| {
+            store.invite_member(channel_id, client_b.id(), proto::ChannelRole::Member, cx)
+        })
+        .await
+        .unwrap();
+
+    // Client B receives a channel invitation notification and responds to the
+    // invitation, accepting it.
+    deterministic.run_until_parked();
+    client_b.notification_store().update(cx_b, |store, cx| {
+        assert_eq!(store.notification_count(), 2);
+        assert_eq!(store.unread_notification_count(), 1);
+
+        let entry = store.notification_at(0).unwrap();
+        assert_eq!(
+            entry.notification,
+            Notification::ChannelInvitation {
+                channel_id,
+                channel_name: "the-channel".to_string(),
+                inviter_id: client_a.id()
+            }
+        );
+        assert!(!entry.is_read);
+
+        store.respond_to_notification(entry.notification.clone(), true, cx);
+    });
+
+    // Client B sees the notification is now read, and that they responded.
+    deterministic.run_until_parked();
+    client_b.notification_store().read_with(cx_b, |store, _| {
+        assert_eq!(store.notification_count(), 2);
+        assert_eq!(store.unread_notification_count(), 0);
+
+        let entry = store.notification_at(0).unwrap();
+        assert!(entry.is_read);
+        assert_eq!(entry.response, Some(true));
+    });
+}