From 69b12f4e331d191cad886568caff85c7c8230a2b Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Wed, 23 Oct 2024 12:34:51 -0400
Subject: [PATCH] semantic_index: Disable embeddings index for non-staff
 (#19618)

This PR disables the embeddings index for non-staff users.

Release Notes:

- N/A
---
 crates/semantic_index/src/embedding_index.rs | 12 ++++++++++++
 crates/semantic_index/src/semantic_index.rs  |  5 +++++
 2 files changed, 17 insertions(+)

diff --git a/crates/semantic_index/src/embedding_index.rs b/crates/semantic_index/src/embedding_index.rs
index a418f60319..0913124341 100644
--- a/crates/semantic_index/src/embedding_index.rs
+++ b/crates/semantic_index/src/embedding_index.rs
@@ -5,6 +5,7 @@ use crate::{
 };
 use anyhow::{anyhow, Context as _, Result};
 use collections::Bound;
+use feature_flags::FeatureFlagAppExt;
 use fs::Fs;
 use futures::stream::StreamExt;
 use futures_batch::ChunksTimeoutStreamExt;
@@ -15,6 +16,7 @@ use log;
 use project::{Entry, UpdatedEntriesSet, Worktree};
 use serde::{Deserialize, Serialize};
 use smol::channel;
+use smol::future::FutureExt;
 use std::{
     cmp::Ordering,
     future::Future,
@@ -65,6 +67,10 @@ impl EmbeddingIndex {
         &self,
         cx: &AppContext,
     ) -> impl Future<Output = Result<()>> {
+        if !cx.is_staff() {
+            return async move { Ok(()) }.boxed();
+        }
+
         let worktree = self.worktree.read(cx).snapshot();
         let worktree_abs_path = worktree.abs_path().clone();
         let scan = self.scan_entries(worktree, cx);
@@ -75,6 +81,7 @@ impl EmbeddingIndex {
             futures::try_join!(scan.task, chunk.task, embed.task, persist)?;
             Ok(())
         }
+        .boxed()
     }
 
     pub fn index_updated_entries(
@@ -82,6 +89,10 @@ impl EmbeddingIndex {
         updated_entries: UpdatedEntriesSet,
         cx: &AppContext,
     ) -> impl Future<Output = Result<()>> {
+        if !cx.is_staff() {
+            return async move { Ok(()) }.boxed();
+        }
+
         let worktree = self.worktree.read(cx).snapshot();
         let worktree_abs_path = worktree.abs_path().clone();
         let scan = self.scan_updated_entries(worktree, updated_entries.clone(), cx);
@@ -92,6 +103,7 @@ impl EmbeddingIndex {
             futures::try_join!(scan.task, chunk.task, embed.task, persist)?;
             Ok(())
         }
+        .boxed()
     }
 
     fn scan_entries(&self, worktree: Snapshot, cx: &AppContext) -> ScanEntries {
diff --git a/crates/semantic_index/src/semantic_index.rs b/crates/semantic_index/src/semantic_index.rs
index 332b4271a0..a18f1f9dbe 100644
--- a/crates/semantic_index/src/semantic_index.rs
+++ b/crates/semantic_index/src/semantic_index.rs
@@ -336,6 +336,11 @@ mod tests {
 
         init_test(cx);
 
+        cx.update(|cx| {
+            // This functionality is staff-flagged.
+            cx.update_flags(true, vec![]);
+        });
+
         let temp_dir = tempfile::tempdir().unwrap();
 
         let mut semantic_index = SemanticDb::new(