defer restoration until initialization is done (#12279)

## Description

Defer restoration of persisted indicies until initailization is done for
remote daemons so we read the correct codebase indexing limit

<!-- Please remember to add your design buddy onto the PR for review, if
it contains any UI changes! -->

## Linked Issue

N/A

<!--
Link the GitHub issue this PR addresses. Before opening this PR, please
confirm:
-->

- [ ] The linked issue is labeled `ready-to-spec` or
`ready-to-implement`.
- [ ] Where appropriate, screenshots or a short video of the
implementation are included below (especially for user-visible or UI
changes).

## Testing

Previously, on startup, these would be erroring. Now they're all
correctly green

![Screenshot 2026-06-05 at
2.57.32 PM.png](https://app.graphite.com/user-attachments/assets/607eb0ee-4220-4fef-a4b4-0f090cb5585a.png)



<!--
How did you test this change? What automated tests did you add? If you
didn't add any new tests, what's your justification for not adding any?

Manual testing is required for changes that can be manually tested, and
almost all changes can be manually tested. If your change can be
manually tested, please include screenshots or a screen recording that
show it working end to end.

You can run the app locally using `./script/run` - see WARP.md for more
details on how to get set up.
-->

- [x] I have manually tested my changes locally with `./script/run`

### Screenshots / Videos

<!-- Attach screenshots or a short video demonstrating the change, where
appropriate. Remove this section if it is not relevant to your PR. -->

## Agent Mode

- [x] Warp Agent Mode - This PR was created via Warp's AI Agent Mode

<!--
## Changelog Entries for Stable

The entries below will be used when constructing a soft-copy of the
stable release changelog. Leave blank or remove the lines if no entry in
the stable changelog is needed. Entries should be on the same line,
without the `{{` `}}` brackets. You can use multiple lines, even of the
same type. The valid suffixes are:

- NEW-FEATURE: for new, relatively sizable features. Features listed
here will likely have docs / social media posts / marketing launches
associated with them, so use sparingly.
- IMPROVEMENT: for new functionality of existing features.
- BUG-FIX: for fixes related to known bugs or regressions.
- IMAGE: the image specified by the URL (hosted on GCP) will be added to
Dev & Preview releases. For Stable releases, see the pinned doc in the
#release Slack channel.
- OZ: Oz-related updates. Use `CHANGELOG-OZ`. At most 4 Oz updates are
shown in-app per release.
- NONE: Explicitly opt out of changelog inclusion. Use `CHANGELOG-NONE`
for PRs that should never appear in the changelog (e.g. refactors,
internal tooling, CI changes). This prevents the changelog agent from
inferring an entry.

CHANGELOG-NEW-FEATURE: {{text goes here...}}
CHANGELOG-IMPROVEMENT: {{text goes here...}}
CHANGELOG-BUG-FIX: {{text goes here...}}
CHANGELOG-BUG-FIX: {{more text goes here...}}
CHANGELOG-IMAGE: {{GCP-hosted URL goes here...}}
CHANGELOG-OZ: {{text goes here...}}
CHANGELOG-NONE
-->

Author: moirahuang

app/src/lib.rs

@@ -1997,14 +1997,17 @@ pub(crate) fn initialize_app(
         };
 
         let codebase_limits = AIRequestUsageModel::as_ref(ctx).codebase_context_limits();
-        let codebase_index_config = CodebaseIndexManagerConfig::new(
+        let mut codebase_index_config = CodebaseIndexManagerConfig::new(
             indices_to_restore,
             codebase_limits.max_indices_allowed,
             codebase_limits.max_files_per_repo,
             codebase_limits.embedding_generation_batch_size,
             server_api_provider.as_ref(ctx).get(),
             launch_mode.supports_indexing(),
         );
+        if matches!(launch_mode, LaunchMode::RemoteServerDaemon { .. }) {
+            codebase_index_config = codebase_index_config.defer_persisted_index_restore();
+        }
         #[cfg(feature = "local_fs")]
         if let Some(snapshot_storage) = daemon_codebase_index_snapshot_storage(launch_mode) {
             return CodebaseIndexManager::new_with_snapshot_storage(

app/src/remote_server/server_model.rs

@@ -1418,6 +1418,9 @@ impl ServerModel {
         log::info!("Handling Initialize (request_id={request_id})");
         self.apply_initialize_auth(&msg);
         Self::apply_codebase_index_limits(msg.codebase_index_limits.as_ref(), ctx);
+        CodebaseIndexManager::handle(ctx).update(ctx, |manager, ctx| {
+            manager.start_persisted_index_restore(ctx);
+        });
 
         // Update crash reporting based on client-supplied preferences.
         #[cfg(feature = "crash_reporting")]

crates/ai/src/index/full_source_code_embedding/manager.rs

@@ -253,6 +253,7 @@ pub struct CodebaseIndexManagerConfig {
     embedding_generation_batch_size: usize,
     store_client: Arc<dyn StoreClient>,
     indexing_enabled: bool,
+    restore_persisted_indices_on_startup: bool,
 }
 
 impl CodebaseIndexManagerConfig {
@@ -271,8 +272,14 @@ impl CodebaseIndexManagerConfig {
             embedding_generation_batch_size,
             store_client,
             indexing_enabled,
+            restore_persisted_indices_on_startup: true,
         }
     }
+
+    pub fn defer_persisted_index_restore(mut self) -> Self {
+        self.restore_persisted_indices_on_startup = false;
+        self
+    }
 }
 
 /// Manager for the codebase index states across the app.
@@ -361,6 +368,7 @@ impl CodebaseIndexManager {
             embedding_generation_batch_size,
             store_client,
             indexing_enabled,
+            restore_persisted_indices_on_startup,
         } = config;
         cfg_if::cfg_if! {
             if #[cfg(feature = "local_fs")] {
@@ -416,7 +424,8 @@ impl CodebaseIndexManager {
         }
 
         // For the moment, we've decided to load all snapshots regardless of the index count.
-        let build_queue = BuildQueue::new_with_persisted(valid_metadata);
+        let build_queue =
+            BuildQueue::new_with_persisted(valid_metadata, restore_persisted_indices_on_startup);
 
         let mut me = Self {
             codebase_indices: HashMap::new(),
@@ -433,10 +442,7 @@ impl CodebaseIndexManager {
             snapshot_storage,
         };
 
-        // Start building the first index in the queue.
-        if let Some(next_repo) = me.build_queue.pick_next_sync() {
-            me.build_and_sync_codebase_index(BuildSource::FromPersistedMetadata(next_repo), ctx);
-        }
+        me.start_next_queued_index(ctx);
 
         me
     }
@@ -808,6 +814,15 @@ impl CodebaseIndexManager {
         self.indexing_enabled
     }
 
+    pub fn start_persisted_index_restore(&mut self, ctx: &mut ModelContext<Self>) {
+        if !self.is_indexing_enabled() {
+            return;
+        }
+        if self.build_queue.start() {
+            self.start_next_queued_index(ctx);
+        }
+    }
+
     pub fn index_directory(&mut self, directory: PathBuf, ctx: &mut ModelContext<Self>) -> bool {
         if !self.is_indexing_enabled() {
             return false;
@@ -1100,7 +1115,10 @@ impl CodebaseIndexManager {
         let Ok(_) = self.get_codebase_index_internal(finished_repo) else {
             return;
         };
+        self.start_next_queued_index(ctx);
+    }
 
+    fn start_next_queued_index(&mut self, ctx: &mut ModelContext<Self>) {
         if let Some(next_repo) = self.build_queue.pick_next_sync() {
             self.build_and_sync_codebase_index(BuildSource::FromPersistedMetadata(next_repo), ctx);
         }

crates/ai/src/index/full_source_code_embedding/manager_tests.rs

@@ -1,6 +1,10 @@
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
+#[cfg(feature = "local_fs")]
+use chrono::Utc;
+#[cfg(feature = "local_fs")]
+use repo_metadata::DirectoryWatcher;
 use warpui_core::App;
 
 use super::{
@@ -174,6 +178,74 @@ fn initializes_with_injected_snapshot_storage_when_configured() {
     });
 }
 
+#[test]
+fn persisted_index_restore_starts_on_startup_by_default() {
+    App::test((), |app| async move {
+        let manager = app.add_singleton_model(|ctx| {
+            CodebaseIndexManager::new(
+                Vec::new(),
+                Some(1),
+                1000,
+                32,
+                Arc::new(MockStoreClient),
+                true,
+                ctx,
+            )
+        });
+
+        manager.read(&app, |manager, _| {
+            assert!(manager.build_queue.is_running());
+        });
+    });
+}
+
+#[test]
+#[cfg(feature = "local_fs")]
+fn deferred_persisted_index_restore_starts_once() {
+    App::test((), |mut app| async move {
+        app.add_singleton_model(DirectoryWatcher::new);
+
+        let snapshot_dir = tempfile::tempdir().unwrap();
+        let storage = SnapshotStorage::from_dir(snapshot_dir.path().join("daemon")).unwrap();
+        let first_repo = tempfile::tempdir().unwrap();
+        let second_repo = tempfile::tempdir().unwrap();
+        let mut first_metadata = workspace_metadata(first_repo.path());
+        first_metadata.modified_ts = Some(Utc::now());
+        let mut second_metadata = workspace_metadata(second_repo.path());
+        second_metadata.modified_ts = Some(Utc::now());
+        std::fs::write(storage.snapshot_path(first_repo.path()), b"snapshot").unwrap();
+        std::fs::write(storage.snapshot_path(second_repo.path()), b"snapshot").unwrap();
+
+        let manager = app.add_singleton_model(|ctx| {
+            CodebaseIndexManager::new_with_snapshot_storage(
+                CodebaseIndexManagerConfig::new(
+                    vec![first_metadata, second_metadata],
+                    Some(2),
+                    1000,
+                    32,
+                    Arc::new(MockStoreClient),
+                    true,
+                )
+                .defer_persisted_index_restore(),
+                Some(storage),
+                ctx,
+            )
+        });
+
+        manager.update(&mut app, |manager, ctx| {
+            assert!(!manager.build_queue.is_running());
+            assert_eq!(manager.build_queue.queued_metadata().into_iter().count(), 2);
+
+            manager.start_persisted_index_restore(ctx);
+            assert!(manager.build_queue.is_running());
+            assert_eq!(manager.build_queue.queued_metadata().into_iter().count(), 1);
+
+            manager.start_persisted_index_restore(ctx);
+            assert_eq!(manager.build_queue.queued_metadata().into_iter().count(), 1);
+        });
+    });
+}
+
 #[test]
 fn can_create_new_indices_honors_max_limit_when_enabled() {
     App::test((), |mut app| async move {

crates/ai/src/index/full_source_code_embedding/priority_queue.rs

@@ -18,6 +18,14 @@ struct QueueEntry {
     metadata: WorkspaceMetadata,
 }
 
+/// Controls whether queued builds may be consumed.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+enum BuildQueueState {
+    Paused,
+    #[default]
+    Running,
+}
+
 impl Hash for QueueEntry {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.metadata.path.hash(state);
@@ -35,6 +43,7 @@ impl Eq for QueueEntry {}
 #[derive(Debug, Default)]
 pub(super) struct BuildQueue {
     queue: PriorityQueue<QueueEntry, Priority>,
+    state: BuildQueueState,
 }
 
 impl BuildQueue {
@@ -46,20 +55,46 @@ impl BuildQueue {
         self.queue.iter().map(|(entry, _)| entry.metadata.clone())
     }
 
-    pub(super) fn new_with_persisted(snapshots_to_load: Vec<WorkspaceMetadata>) -> Self {
+    pub(super) fn new_with_persisted(
+        snapshots_to_load: Vec<WorkspaceMetadata>,
+        start_immediately: bool,
+    ) -> Self {
         let mut queue = PriorityQueue::new();
         queue.extend(
             snapshots_to_load
                 .into_iter()
                 .sorted_by(WorkspaceMetadata::most_recently_touched)
                 .map(|entry| (QueueEntry { metadata: entry }, Priority::PersistedSnapshot)),
         );
+        let state = if start_immediately {
+            BuildQueueState::Running
+        } else {
+            BuildQueueState::Paused
+        };
+
+        Self { queue, state }
+    }
+
+    pub(super) fn is_running(&self) -> bool {
+        self.state == BuildQueueState::Running
+    }
 
-        Self { queue }
+    /// Starts consuming queued builds. Returns whether the queue transitioned to running.
+    pub(super) fn start(&mut self) -> bool {
+        match self.state {
+            BuildQueueState::Paused => {
+                self.state = BuildQueueState::Running;
+                true
+            }
+            BuildQueueState::Running => false,
+        }
     }
 
     /// Pulls the next index root path to sync from the priority queue and returns it.
     pub fn pick_next_sync(&mut self) -> Option<WorkspaceMetadata> {
+        if !self.is_running() {
+            return None;
+        }
         self.queue.pop().map(|(entry, _priority)| entry.metadata)
     }