Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

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

@@ -0,0 +1,50 @@
+use std::collections::HashSet;
+use std::path::PathBuf;
+
+#[derive(Debug, Default, Clone)]
+pub(super) struct ChangedFiles {
+    pub(super) deletions: HashSet<PathBuf>,
+    pub(super) upsertions: HashSet<PathBuf>,
+}
+
+impl ChangedFiles {
+    pub(super) fn is_empty(&self) -> bool {
+        self.deletions.is_empty() && self.upsertions.is_empty()
+    }
+
+    pub(super) fn deletions(&self) -> &HashSet<PathBuf> {
+        &self.deletions
+    }
+
+    /// Merges a subsequent set of file changes into the current set.
+    pub(super) fn merge_subsequent(&mut self, mut subsequent_changes: Self) {
+        for path in subsequent_changes.deletions.drain() {
+            if self.upsertions.contains(&path) {
+                self.upsertions.remove(&path);
+            }
+            self.deletions.insert(path);
+        }
+
+        for path in subsequent_changes.upsertions.drain() {
+            if self.deletions.contains(&path) {
+                self.deletions.remove(&path);
+            }
+            self.upsertions.insert(path);
+        }
+    }
+
+    // Add paths to this changed files set based on whether they currently exist on the file system.
+    pub(super) async fn add_paths(&mut self, paths: impl IntoIterator<Item = PathBuf>) {
+        for path in paths {
+            if path.exists() {
+                self.upsertions.insert(path);
+            } else {
+                self.deletions.insert(path);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "changed_files_test.rs"]
+mod tests;

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

@@ -0,0 +1,251 @@
+use super::*;
+use std::path::PathBuf;
+
+// Helper function to create a PathBuf from a string
+fn pb(path: &str) -> PathBuf {
+    PathBuf::from(path)
+}
+
+#[test]
+fn test_basic_merge_non_conflicting() {
+    // Initial: delete set {a}, upsert set {b}
+    let mut changes1 = ChangedFiles::default();
+    changes1.deletions.insert(pb("a"));
+    changes1.upsertions.insert(pb("b"));
+
+    // Later changes: delete set {c}, upsert set {d}
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("c"));
+    changes2.upsertions.insert(pb("d"));
+
+    // Merge changes
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {a, c}, upsertions {b, d}
+    assert_eq!(changes1.deletions.len(), 2);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert!(changes1.deletions.contains(&pb("c")));
+
+    assert_eq!(changes1.upsertions.len(), 2);
+    assert!(changes1.upsertions.contains(&pb("b")));
+    assert!(changes1.upsertions.contains(&pb("d")));
+}
+
+#[test]
+fn test_delete_then_upsert() {
+    // Initial: delete set {file1}
+    let mut changes1 = ChangedFiles::default();
+    changes1.deletions.insert(pb("file1"));
+
+    // Later changes: upsert set {file1}
+    let mut changes2 = ChangedFiles::default();
+    changes2.upsertions.insert(pb("file1"));
+
+    // Merge changes
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {}, upsertions {file1}
+    assert_eq!(changes1.deletions.len(), 0);
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("file1")));
+}
+
+#[test]
+fn test_upsert_then_delete() {
+    // Initial: upsert set {file1}
+    let mut changes1 = ChangedFiles::default();
+    changes1.upsertions.insert(pb("file1"));
+
+    // Later changes: delete set {file1}
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("file1"));
+
+    // Merge changes
+    changes1.merge_subsequent(changes2);
+
+    // Expected: upsertions {}, deletions {file1}
+    assert_eq!(changes1.upsertions.len(), 0);
+    assert_eq!(changes1.deletions.len(), 1);
+    assert!(changes1.deletions.contains(&pb("file1")));
+}
+
+#[test]
+fn test_delete_then_delete() {
+    // Initial: delete set {file1}
+    let mut changes1 = ChangedFiles::default();
+    changes1.deletions.insert(pb("file1"));
+
+    // Later changes: delete set {file1} again
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("file1"));
+
+    // Merge changes
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {file1}, upsertions {}
+    assert_eq!(changes1.deletions.len(), 1);
+    assert!(changes1.deletions.contains(&pb("file1")));
+    assert_eq!(changes1.upsertions.len(), 0);
+}
+
+#[test]
+fn test_upsert_then_upsert() {
+    // Initial: upsert set {file1}
+    let mut changes1 = ChangedFiles::default();
+    changes1.upsertions.insert(pb("file1"));
+
+    // Later changes: upsert set {file1} again
+    let mut changes2 = ChangedFiles::default();
+    changes2.upsertions.insert(pb("file1"));
+
+    // Merge changes
+    changes1.merge_subsequent(changes2);
+
+    // Expected: upsertions {file1}, deletions {}
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("file1")));
+    assert_eq!(changes1.deletions.len(), 0);
+}
+
+#[test]
+fn test_empty_sets() {
+    // Test 1: Empty merged into populated
+    let mut changes1 = ChangedFiles::default();
+    changes1.upsertions.insert(pb("file1"));
+    changes1.deletions.insert(pb("file2"));
+
+    let changes2 = ChangedFiles::default();
+
+    changes1.merge_subsequent(changes2);
+
+    // Should remain unchanged
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("file1")));
+    assert_eq!(changes1.deletions.len(), 1);
+    assert!(changes1.deletions.contains(&pb("file2")));
+
+    // Test 2: Populated merged into empty
+    let mut changes3 = ChangedFiles::default();
+
+    let mut changes4 = ChangedFiles::default();
+    changes4.upsertions.insert(pb("file3"));
+    changes4.deletions.insert(pb("file4"));
+
+    changes3.merge_subsequent(changes4);
+
+    // Should take all changes
+    assert_eq!(changes3.upsertions.len(), 1);
+    assert!(changes3.upsertions.contains(&pb("file3")));
+    assert_eq!(changes3.deletions.len(), 1);
+    assert!(changes3.deletions.contains(&pb("file4")));
+}
+
+#[test]
+fn test_multiple_sequential_merges() {
+    // Initial: upsert set {a}, delete set {b}
+    let mut changes1 = ChangedFiles::default();
+    changes1.upsertions.insert(pb("a"));
+    changes1.deletions.insert(pb("b"));
+
+    // First merge: delete a, delete c, upsert b
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("a"));
+    changes2.deletions.insert(pb("c"));
+    changes2.upsertions.insert(pb("b"));
+
+    // Second merge: upsert c, delete d
+    let mut changes3 = ChangedFiles::default();
+    changes3.upsertions.insert(pb("c"));
+    changes3.deletions.insert(pb("d"));
+
+    // Apply first merge
+    changes1.merge_subsequent(changes2);
+
+    // After first merge:
+    // Expected: upsertions {b}, deletions {a, c}
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("b")));
+    assert_eq!(changes1.deletions.len(), 2);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert!(changes1.deletions.contains(&pb("c")));
+
+    // Apply second merge
+    changes1.merge_subsequent(changes3);
+
+    // After second merge:
+    // Expected: upsertions {b, c}, deletions {a, d}
+    assert_eq!(changes1.upsertions.len(), 2);
+    assert!(changes1.upsertions.contains(&pb("b")));
+    assert!(changes1.upsertions.contains(&pb("c")));
+    assert_eq!(changes1.deletions.len(), 2);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert!(changes1.deletions.contains(&pb("d")));
+}
+
+#[test]
+fn test_rename_then_delete() {
+    // Initial: rename {a -> b}
+    let mut changes1 = ChangedFiles::default();
+    changes1.deletions.insert(pb("a"));
+    changes1.upsertions.insert(pb("b"));
+
+    // Later changes: delete {b}
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("b"));
+
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {a, b}, upsertions {}
+    // Note that we don't know whether b had any prior content,
+    // so we can't assume it was a rename.
+    assert_eq!(changes1.deletions.len(), 2);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert!(changes1.deletions.contains(&pb("b")));
+    assert_eq!(changes1.upsertions.len(), 0);
+}
+
+#[test]
+fn test_upsert_then_rename() {
+    // Initial: upsert {a}
+    let mut changes1 = ChangedFiles::default();
+    changes1.upsertions.insert(pb("a"));
+
+    // Later changes: rename {a -> b}
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("a"));
+    changes2.upsertions.insert(pb("b"));
+
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {a}, upsertions {b}
+    // Note that we don't know whether a had any prior content,
+    // so we can't assume it was a rename.
+    assert_eq!(changes1.deletions.len(), 1);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("b")));
+}
+
+#[test]
+fn test_rename_then_rename() {
+    // Initial: rename {a -> b}
+    let mut changes1 = ChangedFiles::default();
+    changes1.deletions.insert(pb("a"));
+    changes1.upsertions.insert(pb("b"));
+
+    // Later changes: rename {b -> c}
+    let mut changes2 = ChangedFiles::default();
+    changes2.deletions.insert(pb("b"));
+    changes2.upsertions.insert(pb("c"));
+
+    changes1.merge_subsequent(changes2);
+
+    // Expected: deletions {a, b}, upsertions {c}
+    // Note that we don't know whether b had any prior content,
+    // so we can't assume it was a rename.
+    assert_eq!(changes1.deletions.len(), 2);
+    assert!(changes1.deletions.contains(&pb("a")));
+    assert!(changes1.deletions.contains(&pb("b")));
+    assert_eq!(changes1.upsertions.len(), 1);
+    assert!(changes1.upsertions.contains(&pb("c")));
+}

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

@@ -0,0 +1,109 @@
+use std::path::Path;
+
+use string_offset::ByteOffset;
+
+mod naive;
+#[cfg(not(target_family = "wasm"))]
+mod semantic;
+
+/// Number of lines per chunk when chunking naively. While there's no guarantee
+/// that this is below the token limit of the embedding model used on the server,
+/// this should give us more than enough buffer.
+const LINES_PER_CHUNK: usize = 200;
+
+/// The average number of characters per line.
+const AVG_CHAR_PER_LINE: usize = 60;
+
+/// Compute the max byte per chunk based on the average number of characters per line. We assume code is mostly ASCII,
+/// which is why this max chunk makes sense even if we're using bytes instead of characters as our unit of chunking.
+const MAX_BYTES_PER_CHUNK: usize = LINES_PER_CHUNK * AVG_CHAR_PER_LINE;
+
+/// A code fragment with line range information.
+#[derive(Debug, Clone)]
+pub struct Fragment<'a> {
+    /// The content of the fragment.
+    pub content: &'a str,
+    /// Start line number (inclusive).
+    pub start_line: usize,
+    /// End line number (inclusive).
+    pub end_line: usize,
+    /// The start byte index of the fragment in the original source code.
+    pub start_byte_index: ByteOffset,
+    /// The end byte index of the fragment (exclusive) in the original source code.
+    pub end_byte_index: ByteOffset,
+    /// File path of the fragment.
+    pub file_path: &'a Path,
+}
+
+impl<'a> Fragment<'a> {
+    fn size(&self) -> usize {
+        self.content.len()
+    }
+
+    fn append(&mut self, other: &Fragment<'a>, content: &'a str) {
+        self.end_line = other.end_line;
+        self.end_byte_index = other.end_byte_index;
+        self.content = &content[self.start_byte_index.as_usize()..other.end_byte_index.as_usize()];
+    }
+}
+
+/// Coalesce small fragments into larger ones that still respect the `max_bytes_per_chunk`.
+/// Treesitter often produces small fragments that splits function names from the actual function body,
+/// we iterate in reverse to coalesce these chunks into fragments that are more meaningful.
+fn coalesce_fragments<'a>(
+    fragments: impl DoubleEndedIterator<Item = Fragment<'a>>,
+    code: &'a str,
+    max_bytes_per_chunk: usize,
+) -> Vec<Fragment<'a>> {
+    fragments
+        .rev()
+        .fold(
+            Vec::new(),
+            |mut acc: Vec<Fragment<'a>>, mut fragment| match acc.last_mut() {
+                Some(last_item) => {
+                    let new_fragment_size = code
+                        [fragment.start_byte_index.as_usize()..last_item.end_byte_index.as_usize()]
+                        .len();
+                    if new_fragment_size <= max_bytes_per_chunk {
+                        fragment.append(last_item, code);
+                        *last_item = fragment;
+                    } else {
+                        acc.push(fragment);
+                    }
+                    acc
+                }
+                None => {
+                    acc.push(fragment);
+                    acc
+                }
+            },
+        )
+        .into_iter()
+        .rev()
+        .collect()
+}
+
+/// Chunks code into an ordered list of fragments.
+///
+/// The code is chunked "semantically" using treesitter.
+/// If we are unable to generate semantic chunks for any reason, fragments are naively chunked by
+/// lines.
+pub fn chunk_code<'a>(code: &'a str, path: &'a Path) -> Vec<Fragment<'a>> {
+    if let Some(fragments) = try_chunk_code_semantically(code, path) {
+        return fragments;
+    }
+    naive::chunk_code(code, path, MAX_BYTES_PER_CHUNK, LINES_PER_CHUNK)
+}
+
+/// Attempts to chunk code semantically, returning [`None`] if the code
+/// could not be chunked for any reason.
+#[cfg(not(target_family = "wasm"))]
+fn try_chunk_code_semantically<'a>(code: &'a str, path: &'a Path) -> Option<Vec<Fragment<'a>>> {
+    let language = languages::language_by_filename(path)?;
+    semantic::chunk_code(code, path, MAX_BYTES_PER_CHUNK, &language.grammar).ok()
+}
+
+#[cfg(target_family = "wasm")]
+fn try_chunk_code_semantically<'a>(_code: &'a str, _path: &'a Path) -> Option<Vec<Fragment<'a>>> {
+    None
+}

crates/ai/src/index/full_source_code_embedding/chunker/naive.rs

@@ -0,0 +1,108 @@
+use crate::index::full_source_code_embedding::chunker::{coalesce_fragments, Fragment};
+use itertools::Itertools;
+use line_span::{LineSpan, LineSpans};
+use std::path::Path;
+
+/// Chunks the given file into [`Fragment`]s. Each chunk is at most `num_lines_per_chunk` lines long, and contains at most `max_bytes_per_chunk` bytes.
+pub(super) fn chunk_code<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    num_lines_per_chunk: usize,
+) -> Vec<Fragment<'a>> {
+    let lines = code.line_spans().enumerate().collect_vec();
+    let chunks = lines.chunks(num_lines_per_chunk);
+
+    chunks
+        .into_iter()
+        .flat_map(|chunk| {
+            let (start_line, start_range) = chunk[0];
+            let (end_line, end_range) =
+                chunk.last().expect("Chunks must have at least one element");
+
+            if (end_range.end() - start_range.start()) > max_bytes_per_chunk {
+                let chunked_fragments = chunk.iter().flat_map(|(line, line_span)| {
+                    chunk_line_by_bytes(code, path, max_bytes_per_chunk, *line, line_span)
+                });
+
+                return coalesce_fragments(chunked_fragments, code, max_bytes_per_chunk);
+            }
+
+            vec![Fragment {
+                content: &code[start_range.start()..end_range.end()],
+                start_line,
+                end_line: *end_line,
+                file_path: path,
+                start_byte_index: start_range.start().into(),
+                end_byte_index: end_range.end().into(),
+            }]
+        })
+        .collect()
+}
+
+/// Chunks the line represented by `line_span` into multiple fragments if it exceeds `max_bytes_per_chunk`.
+fn chunk_line_by_bytes<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    line_number: usize,
+    line_span: &LineSpan<'a>,
+) -> Vec<Fragment<'a>> {
+    let line_start = line_span.start();
+    let line_end = line_span.end();
+    let line_content = &code[line_start..line_end];
+    let line_length = line_end - line_start;
+
+    // If the line is smaller than max_bytes_per_chunk, return it as a single fragment
+    if line_length <= max_bytes_per_chunk {
+        return vec![Fragment {
+            content: line_content,
+            start_line: line_number,
+            end_line: line_number,
+            file_path: path,
+            start_byte_index: line_start.into(),
+            end_byte_index: line_end.into(),
+        }];
+    }
+
+    // Otherwise, split the line into multiple fragments
+    let mut fragments = Vec::new();
+    let mut current_start = line_start;
+
+    while current_start < line_end {
+        let remaining_bytes = line_end - current_start;
+        let chunk_size = std::cmp::min(remaining_bytes, max_bytes_per_chunk);
+        let mut chunk_end = current_start + chunk_size;
+
+        // Ensure chunk_end is on a UTF-8 character boundary
+        while chunk_end > current_start && !code.is_char_boundary(chunk_end) {
+            chunk_end -= 1;
+        }
+
+        // If we couldn't find a valid boundary within reasonable distance,
+        // move forward to the next character boundary instead
+        if chunk_end <= current_start {
+            chunk_end = current_start + chunk_size;
+            while chunk_end < line_end && !code.is_char_boundary(chunk_end) {
+                chunk_end += 1;
+            }
+        }
+
+        fragments.push(Fragment {
+            content: &code[current_start..chunk_end],
+            start_line: line_number,
+            end_line: line_number,
+            file_path: path,
+            start_byte_index: current_start.into(),
+            end_byte_index: chunk_end.into(),
+        });
+
+        current_start = chunk_end;
+    }
+
+    fragments
+}
+
+#[cfg(test)]
+#[path = "naive_tests.rs"]
+mod tests;

crates/ai/src/index/full_source_code_embedding/chunker/naive_tests.rs

@@ -0,0 +1,354 @@
+use super::*;
+use std::path::Path;
+
+#[test]
+fn test_chunker() {
+    let code = "This is some text content\nthat should be chunked\nusing the naive chunker\nbecause the language isn't recognized.";
+    let path = Path::new("test_file.xyz");
+
+    let max_lines = 1;
+    let fragments = chunk_code(code, path, 10000, max_lines);
+
+    assert!(!fragments.is_empty(), "Expected at least one fragment");
+
+    assert_eq!(fragments.len(), code.lines().count());
+    for (idx, line) in code.lines().enumerate() {
+        assert_eq!(fragments[idx].content, line);
+        assert_eq!(fragments[idx].start_line, idx);
+        assert_eq!(fragments[idx].end_line, idx);
+    }
+}
+
+#[test]
+fn test_chunker_large_chunk() {
+    let code = "This is some text content\nthat should be chunked\nusing the naive chunker\nbecause the language isn't recognized.";
+    let path = Path::new("test_file.xyz");
+
+    let fragments = chunk_code(code, path, 10000, 100);
+
+    // We should have only one fragment
+    assert_eq!(fragments.len(), 1);
+
+    assert_eq!(fragments[0].content, code);
+    assert_eq!(fragments[0].start_line, 0);
+    assert_eq!(fragments[0].end_line, code.lines().count() - 1);
+}
+
+#[test]
+fn test_chunker_max_bytes() {
+    // Create a string with known byte size - each line is exactly 20 bytes including newline
+    let code = "line1\nline2\nline3\nline4abcdefghijklmnopqrstuvwxyz";
+    let path = Path::new("test_file.xyz");
+
+    // Set max_bytes_per_chunk to 25 bytes to force multiple chunks for the last line (which is 30 bytes).
+    let max_bytes_per_chunk = 25;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that no chunk exceeds the max_bytes_per_chunk limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.trim().len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk
+        );
+    }
+
+    // The first fragment contains all of the lines except the last one.
+    assert_eq!(fragments[0].content, "line1\nline2\nline3");
+
+    // The last two fragments contains the contents of the line line.
+    assert_eq!(fragments[1].content, "line4abcdefghijklmnopqrst");
+    assert_eq!(fragments[2].content, "uvwxyz");
+
+    // Verify that the chunks together contain all the original content
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    // Ignore any newlines when doing comparisons--the chunker may drop newlines at fragment boundaries
+    // and that's not necessary for testing the correctness of the naive chunker.
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_emoji_chunking() {
+    // Test with emojis (4-byte UTF-8 characters) to ensure byte boundaries are respected
+    let code = "Hello 🦀 Rust\nWorld 🌍 Test\n🚀 Rocket 🎯 Target";
+    let path = Path::new("test_emoji.txt");
+
+    // Set a small max_bytes_per_chunk to force splitting through emoji characters
+    let max_bytes_per_chunk = 15; // This will force splits in the middle of emoji sequences
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that no chunk exceeds the max_bytes_per_chunk limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes. Content: '{}'",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk,
+            fragment.content
+        );
+    }
+
+    // Verify that all fragments contain valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.is_ascii() || std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content
+        );
+    }
+
+    // Verify that reassembled content matches original (ignoring newlines)
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_accented_characters() {
+    // Test with accented characters (2-byte UTF-8)
+    let code = "Café résumé naïve\nÉlève découvrir\nMañana piñata";
+    let path = Path::new("test_accents.txt");
+
+    // Set max_bytes_per_chunk to force splitting through accented characters
+    let max_bytes_per_chunk = 10;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that all fragments contain valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify that reassembled content matches original (ignoring newlines)
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_mixed_characters() {
+    // Test with a mix of ASCII, 2-byte, 3-byte, and 4-byte UTF-8 characters
+    let code = "ASCII text 中文 🦀 résumé ℘ math symbols";
+    let path = Path::new("test_mixed.txt");
+
+    // Set a small chunk size to force many splits
+    let max_bytes_per_chunk = 8;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that all fragments contain valid UTF-8 and don't exceed size limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk
+        );
+
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify that reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_boundary_edge_cases() {
+    // Test edge case where chunk boundary falls exactly on a multi-byte character
+    let code = "ab🦀cd"; // 'ab' (2 bytes) + '🦀' (4 bytes) + 'cd' (2 bytes) = 8 bytes total
+    let path = Path::new("test_edge.txt");
+
+    // Set chunk size to 3 bytes, which would split in the middle of the emoji without our fix
+    let max_bytes_per_chunk = 3;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Should have multiple fragments
+    assert!(fragments.len() >= 2, "Expected at least 2 fragments");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_single_multibyte_character() {
+    // Test with a single multi-byte character that's larger than chunk size
+    let code = "🦀"; // 4-byte emoji
+    let path = Path::new("test_single.txt");
+
+    // Set chunk size smaller than the character
+    let max_bytes_per_chunk = 2;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Should have exactly one fragment (can't split a single character)
+    assert_eq!(fragments.len(), 1, "Should have exactly one fragment");
+
+    // The fragment should contain the complete character
+    assert_eq!(fragments[0].content, code);
+
+    // Verify it's valid UTF-8
+    assert!(
+        std::str::from_utf8(fragments[0].content.as_bytes()).is_ok(),
+        "Fragment contains invalid UTF-8"
+    );
+}
+
+#[test]
+fn test_utf8_line_endings_with_multibyte() {
+    // Test multi-byte characters at line boundaries
+    let code = "Hello🌍\nWorld🦀\nTest🎯";
+    let path = Path::new("test_lines.txt");
+
+    let max_bytes_per_chunk = 10;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1); // 1 line per chunk
+
+    // Should have 3 fragments (one per line)
+    assert_eq!(fragments.len(), 3, "Should have 3 fragments for 3 lines");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify line numbers are correct
+    assert_eq!(fragments[0].start_line, 0);
+    assert_eq!(fragments[0].end_line, 0);
+    assert_eq!(fragments[1].start_line, 1);
+    assert_eq!(fragments[1].end_line, 1);
+    assert_eq!(fragments[2].start_line, 2);
+    assert_eq!(fragments[2].end_line, 2);
+}
+
+#[test]
+fn test_panic_regression_byte_boundary() {
+    // This is a regression test for the "byte index is not a char boundary" panic.
+    // Before the fix, this would panic when trying to slice at byte index 3,
+    // which is in the middle of the 4-byte emoji '🦀'.
+    let code = "Hi🦀Test";
+    let path = Path::new("test_panic.txt");
+
+    // This chunk size would cause the original code to panic
+    let max_bytes_per_chunk = 3;
+
+    // This should not panic
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we get valid fragments
+    assert!(!fragments.is_empty(), "Should have at least one fragment");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}

crates/ai/src/index/full_source_code_embedding/chunker/semantic.rs

@@ -0,0 +1,170 @@
+use std::path::Path;
+
+use arborium::tree_sitter::{Language, Node, Parser, TreeCursor};
+use itertools::Itertools;
+
+use super::{coalesce_fragments, Fragment};
+
+/// Maximum depth for recursive tree traversal to prevent infinite recursion
+/// or excessive depth in malformed/deeply nested code.
+const MAX_TRAVERSAL_DEPTH: usize = 200;
+
+/// Chunks code into an ordered list of fragments, where each fragment is at most
+/// `max_bytes_per_chunk` bytes.
+pub(super) fn chunk_code<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    language: &Language,
+) -> anyhow::Result<Vec<Fragment<'a>>> {
+    // Wrap this in a block to ensure the treesitter Parser / Tree are dropped
+    // after creating the fragments.
+    let fragments = {
+        let mut parser = Parser::new();
+        parser.set_language(language)?;
+
+        let tree = parser
+            .parse(code, None /* old_tree */)
+            .ok_or_else(|| anyhow::anyhow!("Failed to parse code"))?;
+
+        let mut cursor = tree.walk();
+
+        let nodes = split_node(
+            tree.root_node(),
+            code,
+            max_bytes_per_chunk,
+            path,
+            &mut cursor,
+            0, // initial depth
+        )?;
+
+        coalesce_fragments(nodes.into_iter(), code, max_bytes_per_chunk)
+    };
+
+    // Release extra unused memory from malloc to the system.  For some
+    // reason, the memory obtained by the allocator is often not released
+    // back to the OS after we're done with it, resulting in high memory
+    // usage (from the perspective of the OS, though not from the perspective
+    // of the allocator).
+    //
+    // See: https://github.com/tree-sitter/tree-sitter/issues/3129
+    #[cfg(all(target_os = "linux", target_env = "gnu", not(feature = "jemalloc")))]
+    unsafe {
+        nix::libc::malloc_trim(0);
+    }
+
+    Ok(fragments)
+}
+
+/// Splits a [`Node`] into a series of [`Fragment`]s that are at most `max_bytes_per_chunk` bytes.
+fn split_node<'a, 'b>(
+    node: Node<'b>,
+    code: &'a str,
+    max_bytes_per_chunk: usize,
+    path: &'a Path,
+    cursor: &mut TreeCursor<'b>,
+    depth: usize,
+) -> anyhow::Result<Vec<Fragment<'a>>> {
+    // Check if we've exceeded the maximum traversal depth
+    if depth > MAX_TRAVERSAL_DEPTH {
+        return Err(anyhow::anyhow!(
+            "Maximum traversal depth {} exceeded, falling back to naive chunking",
+            MAX_TRAVERSAL_DEPTH
+        ));
+    }
+
+    let mut current_fragment = Fragment::from_node_start(node, path);
+    let mut fragments = vec![];
+
+    // Collect into a vec to avoid a double mutable borrow with `cursor` when we make
+    // the recursive call below.
+    for child in node.children(cursor).collect_vec() {
+        let child_size = child.end_byte().saturating_sub(child.start_byte());
+
+        // The child is larger than the max chunk size, so we need to split it recursively.
+        if child_size > max_bytes_per_chunk {
+            let mut new_fragment = Fragment::from_node_end(child, path);
+            std::mem::swap(&mut current_fragment, &mut new_fragment);
+            fragments.push(new_fragment);
+
+            fragments.append(&mut split_node(
+                child,
+                code,
+                max_bytes_per_chunk,
+                path,
+                cursor,
+                depth + 1,
+            )?);
+        } else if child_size + current_fragment.size() > max_bytes_per_chunk {
+            // The child would make the current fragment too large, so we finalize the current
+            // fragment and create a new one.
+            fragments.push(current_fragment);
+            current_fragment = Fragment::from_node_start(child, path);
+            current_fragment.append(&Fragment::from_node_end(child, path), code);
+        } else {
+            // The child fits within the current fragment.
+            current_fragment.end_line = child.end_position().row;
+            current_fragment.end_byte_index = child.end_byte().into();
+            current_fragment.content =
+                &code[current_fragment.start_byte_index.as_usize()..child.end_byte()];
+        }
+    }
+
+    fragments.push(current_fragment);
+
+    Ok(fragments)
+}
+
+impl<'a> Fragment<'a> {
+    /// Creates an empty fragment.
+    fn empty() -> Fragment<'a> {
+        Fragment {
+            content: "",
+            start_line: 0,
+            end_line: 0,
+            start_byte_index: 0.into(),
+            end_byte_index: 0.into(),
+            file_path: Path::new(""),
+        }
+    }
+
+    /// Creates a fragment comprised solely of the start of the given node.
+    fn from_node_start(node: Node<'_>, path: &'a Path) -> Self {
+        Fragment {
+            content: "",
+            start_line: node.start_position().row,
+            end_line: node.start_position().row,
+            start_byte_index: node.start_byte().into(),
+            end_byte_index: node.start_byte().into(),
+            file_path: path,
+        }
+    }
+
+    /// Creates a fragment comprised solely of the end of the given node.
+    fn from_node_end(node: Node<'_>, path: &'a Path) -> Self {
+        Fragment {
+            content: "",
+            start_line: node.end_position().row,
+            end_line: node.end_position().row,
+            start_byte_index: node.end_byte().into(),
+            end_byte_index: node.end_byte().into(),
+            file_path: path,
+        }
+    }
+
+    /// Creates a fragment comprised solely of the end of the given fragment.
+    fn from_fragment_end(fragment: &Fragment<'a>) -> Self {
+        Fragment {
+            content: "",
+            start_line: fragment.end_line,
+            end_line: fragment.end_line,
+            start_byte_index: fragment.end_byte_index,
+            end_byte_index: fragment.end_byte_index,
+            file_path: fragment.file_path,
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "semantic_tests.rs"]
+mod tests;

crates/ai/src/index/full_source_code_embedding/chunker/semantic_tests.rs

@@ -0,0 +1,90 @@
+use std::path::Path;
+
+use languages::language_by_filename;
+
+use super::*;
+
+#[test]
+fn test_basic_rust_chunking() {
+    let source_code = r#"
+#[derive(Debug)]
+struct Rectangle {
+    width: u32,
+    height: u32,
+}
+
+impl Rectangle {
+    fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}
+
+fn main() {
+    let rect1 = Rectangle {
+        width: 30,
+        height: 50,
+    };
+
+    println!(
+        "The area of the rectangle is {} square pixels.",
+        rect1.area()
+    );
+}
+"#;
+
+    let max_chunk_size = 128;
+
+    let chunks = chunk_code(
+        source_code,
+        Path::new("test.rs"),
+        max_chunk_size,
+        &language_by_filename(Path::new("test.rs"))
+            .expect("Rust language must exist")
+            .grammar,
+    )
+    .unwrap();
+
+    assert_eq!(chunks.len(), 4);
+
+    // None of the chunks should exceed the chunk size.
+    for chunk in &chunks {
+        assert!(
+            chunk.content.len() <= max_chunk_size,
+            "Chunk should not exceed max size of {max_chunk_size} but was: {}",
+            chunk.content.len()
+        );
+    }
+
+    assert_eq!(
+        chunks[0].content.trim(),
+        r#"#[derive(Debug)]
+struct Rectangle {
+    width: u32,
+    height: u32,
+}"#
+    );
+    assert_eq!(
+        chunks[1].content.trim(),
+        r#"impl Rectangle {
+    fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}"#
+    );
+    assert_eq!(
+        chunks[2].content.trim(),
+        r#"fn main() {
+    let rect1 = Rectangle {
+        width: 30,
+        height: 50,
+    };"#
+    );
+    assert_eq!(
+        chunks[3].content.trim(),
+        r#"println!(
+        "The area of the rectangle is {} square pixels.",
+        rect1.area()
+    );
+}"#
+    );
+}

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

@@ -0,0 +1,148 @@
+use std::{collections::HashMap, ops::Range, path::PathBuf};
+
+use serde::{Deserialize, Serialize};
+use string_offset::ByteOffset;
+
+use super::merkle_tree::MerkleHash;
+use crate::index::full_source_code_embedding::chunker::Fragment;
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct FragmentLocation {
+    pub start_line: usize,
+    /// End line number (inclusive).
+    pub end_line: usize,
+    /// The range of byte indices into the original source string for this fragment.
+    pub byte_range: Range<ByteOffset>,
+}
+
+/// Fragment metadata that we persist in the tree. This helps us map from a leaf merkle node
+/// to the actual content on user's disk.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct FragmentMetadata {
+    /// File path of the fragment.
+    pub absolute_path: PathBuf,
+    /// Location of the fragment within the file.
+    pub location: FragmentLocation,
+}
+
+impl FragmentMetadata {
+    /// Returns the estimated content size in bytes, derived from the stored byte range.
+    pub fn content_byte_size(&self) -> usize {
+        self.location
+            .byte_range
+            .end
+            .as_usize()
+            .saturating_sub(self.location.byte_range.start.as_usize())
+    }
+}
+
+impl From<&Fragment<'_>> for FragmentMetadata {
+    fn from(fragment: &Fragment<'_>) -> Self {
+        FragmentMetadata {
+            absolute_path: PathBuf::from(fragment.file_path),
+            location: FragmentLocation {
+                start_line: fragment.start_line,
+                end_line: fragment.end_line,
+                byte_range: fragment.start_byte_index..fragment.end_byte_index,
+            },
+        }
+    }
+}
+
+pub type LeafToFragmentMetadataMapping = HashMap<MerkleHash, Vec<FragmentMetadata>>;
+
+#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
+pub struct LeafToFragmentMetadata {
+    mapping: LeafToFragmentMetadataMapping,
+}
+
+impl LeafToFragmentMetadata {
+    pub(super) fn empty() -> Self {
+        Self::default()
+    }
+
+    pub(super) fn new(initial_content: LeafToFragmentMetadataUpdates) -> Self {
+        Self::from(initial_content)
+    }
+
+    #[cfg(test)]
+    pub(super) fn new_for_test(content: HashMap<MerkleHash, Vec<FragmentMetadata>>) -> Self {
+        Self { mapping: content }
+    }
+
+    pub(super) fn mapping(&self) -> &LeafToFragmentMetadataMapping {
+        &self.mapping
+    }
+
+    pub fn get<T: AsRef<MerkleHash>>(&self, hash: T) -> Option<&Vec<FragmentMetadata>> {
+        self.mapping.get(hash.as_ref())
+    }
+
+    pub fn apply_update(&mut self, update: LeafToFragmentMetadataUpdates) {
+        let LeafToFragmentMetadataUpdates {
+            to_remove,
+            to_insert,
+        } = update;
+        for (path, hashes) in to_remove {
+            for hash in hashes {
+                let Some(mapping_entry) = self.mapping.get_mut(&hash) else {
+                    continue;
+                };
+                mapping_entry.retain(|metadata| metadata.absolute_path != path);
+                if mapping_entry.is_empty() {
+                    self.mapping.remove(&hash);
+                }
+            }
+        }
+        to_insert.into_iter().for_each(|(hash, metadatas)| {
+            self.mapping.entry(hash).or_default().extend(metadatas);
+        });
+    }
+}
+
+impl From<LeafToFragmentMetadataUpdates> for LeafToFragmentMetadata {
+    fn from(update: LeafToFragmentMetadataUpdates) -> Self {
+        Self {
+            mapping: update.to_insert,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct LeafToFragmentMetadataUpdates {
+    /// Since the same fragment can occur multiple times within the same file,
+    /// the filepath is not enough to uniquely identify a fragment.
+    /// At the moment, we only handle removing entire files, so using the path alone is okay.
+    /// In the future, add the FragmentMetadata to the key to uniquely identify a fragment.
+    pub(super) to_remove: HashMap<PathBuf, Vec<MerkleHash>>,
+    pub(super) to_insert: LeafToFragmentMetadataMapping,
+}
+
+impl LeafToFragmentMetadataUpdates {
+    pub fn empty() -> Self {
+        Self::default()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.to_remove.is_empty() && self.to_insert.is_empty()
+    }
+
+    pub fn merge(&mut self, other: Self) {
+        other.to_remove.into_iter().for_each(|(path, hashes)| {
+            self.to_remove.entry(path).or_default().extend(hashes);
+        });
+        other.to_insert.into_iter().for_each(|(hash, metadata)| {
+            self.to_insert.entry(hash).or_default().extend(metadata);
+        })
+    }
+
+    pub fn insertions(&self) -> &LeafToFragmentMetadataMapping {
+        &self.to_insert
+    }
+}
+
+impl Extend<LeafToFragmentMetadataUpdates> for LeafToFragmentMetadataUpdates {
+    fn extend<T: IntoIterator<Item = LeafToFragmentMetadataUpdates>>(&mut self, iter: T) {
+        iter.into_iter().for_each(|update| self.merge(update));
+    }
+}

crates/ai/src/index/full_source_code_embedding/merkle_tree/hash.rs

@@ -0,0 +1,227 @@
+//! Common types for hashes that identify codebase embedding state.
+
+use generic_array::GenericArray;
+use serde::{Deserialize, Serialize};
+use sha2::{digest::OutputSizeUser, Digest, Sha256};
+use std::{fmt, str::FromStr, sync::Arc};
+
+use crate::index::full_source_code_embedding::chunker::Fragment;
+
+use super::Error;
+
+/// The hash of an *intermediate* node in the [`MerkleTree`].
+///
+/// Unlike [`MerkleHash`], this is guaranteed to be an intermediate node.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct NodeHash(MerkleHash);
+
+impl NodeHash {
+    pub(super) fn new(hash: MerkleHash) -> Self {
+        Self(hash)
+    }
+}
+
+impl fmt::Display for NodeHash {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl FromStr for NodeHash {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self(MerkleHash::from_str(s)?))
+    }
+}
+
+impl From<NodeHash> for warp_graphql::full_source_code_embedding::NodeHash {
+    fn from(value: NodeHash) -> Self {
+        warp_graphql::full_source_code_embedding::NodeHash(value.0.to_string())
+    }
+}
+
+impl TryFrom<warp_graphql::full_source_code_embedding::NodeHash> for NodeHash {
+    type Error = Error;
+
+    fn try_from(
+        value: warp_graphql::full_source_code_embedding::NodeHash,
+    ) -> Result<Self, Self::Error> {
+        Ok(Self(MerkleHash::from_str(&value.0)?))
+    }
+}
+
+impl AsRef<MerkleHash> for NodeHash {
+    fn as_ref(&self) -> &MerkleHash {
+        &self.0
+    }
+}
+
+impl From<&ContentHash> for NodeHash {
+    fn from(value: &ContentHash) -> Self {
+        NodeHash(value.0.to_owned())
+    }
+}
+
+impl From<ContentHash> for NodeHash {
+    fn from(value: ContentHash) -> Self {
+        NodeHash(value.0)
+    }
+}
+
+/// The hash of a fragment (leaf) node in the [`MerkleTree`].
+///
+/// Unlike [`MerkleHash`], this is guaranteed to be a leaf node.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ContentHash(MerkleHash);
+
+impl ContentHash {
+    pub(crate) fn new(hash: MerkleHash) -> Self {
+        Self(hash)
+    }
+
+    pub fn from_content(content: &str) -> Self {
+        Self(MerkleHash::from_bytes(content.as_bytes()))
+    }
+}
+
+impl fmt::Display for ContentHash {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl FromStr for ContentHash {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self(MerkleHash::from_str(s)?))
+    }
+}
+
+impl From<ContentHash> for warp_graphql::full_source_code_embedding::ContentHash {
+    fn from(value: ContentHash) -> Self {
+        warp_graphql::full_source_code_embedding::ContentHash(value.0.to_string())
+    }
+}
+
+impl TryFrom<warp_graphql::full_source_code_embedding::ContentHash> for ContentHash {
+    type Error = Error;
+
+    fn try_from(
+        value: warp_graphql::full_source_code_embedding::ContentHash,
+    ) -> Result<Self, Self::Error> {
+        Ok(Self(MerkleHash::from_str(&value.0)?))
+    }
+}
+
+impl AsRef<MerkleHash> for ContentHash {
+    fn as_ref(&self) -> &MerkleHash {
+        &self.0
+    }
+}
+
+impl AsRef<ContentHash> for ContentHash {
+    fn as_ref(&self) -> &ContentHash {
+        self
+    }
+}
+
+/// A SHA-256 hash for a node in the [`MerkleTree`].
+///
+/// Cloning a `MerkleHash` is cheap, and need not be avoided.
+/// TODO(CODE-399): make this private to the `merkle_tree` module.
+#[derive(Ord, PartialOrd, Eq, PartialEq, Hash, Clone)]
+pub(crate) struct MerkleHash(Arc<GenericArray<u8, <Sha256 as OutputSizeUser>::OutputSize>>);
+
+impl AsRef<MerkleHash> for MerkleHash {
+    fn as_ref(&self) -> &MerkleHash {
+        self
+    }
+}
+
+/// The default serialize prints the hash as a vector of small integers,
+/// which takes up more space (up to 5 characters per byte) and is more
+/// difficult to read.
+/// This custom serialization hex-encodes the bytes in a string instead.
+impl Serialize for MerkleHash {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        let mut buf = [0u8; 64];
+        let hex_str = base16ct::lower::encode_str(&self.0, &mut buf)
+            .expect("Buffer is sufficient for a SHA-256 hash");
+        serializer.serialize_str(hex_str)
+    }
+}
+
+impl<'de> Deserialize<'de> for MerkleHash {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let hex_string = String::deserialize(deserializer)?;
+        MerkleHash::from_str(&hex_string).map_err(serde::de::Error::custom)
+    }
+}
+
+impl MerkleHash {
+    pub(super) fn from_hashes<'a>(iterator: impl Iterator<Item = &'a MerkleHash>) -> Self {
+        let mut hasher = Sha256::new();
+        for hash in iterator {
+            Digest::update(&mut hasher, hash.0.as_slice());
+        }
+
+        Self::from_digest(hasher)
+    }
+
+    pub(crate) fn from_bytes(content_bytes: &[u8]) -> Self {
+        let mut hasher = Sha256::new();
+        Digest::update(&mut hasher, content_bytes);
+
+        Self::from_digest(hasher)
+    }
+
+    pub(super) fn from_fragment(fragment: &Fragment<'_>) -> Self {
+        Self::from_bytes(fragment.content.as_bytes())
+    }
+
+    fn from_digest(digest: Sha256) -> Self {
+        Self(Arc::new(digest.finalize()))
+    }
+}
+
+impl FromStr for MerkleHash {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let mut buf: GenericArray<u8, <Sha256 as OutputSizeUser>::OutputSize> = Default::default();
+        let decoded =
+            base16ct::lower::decode(s.as_bytes(), &mut buf).map_err(Error::InvalidHash)?;
+        if decoded.len() != 32 {
+            return Err(Error::InvalidHash(base16ct::Error::InvalidLength));
+        }
+
+        Ok(MerkleHash(Arc::new(buf)))
+    }
+}
+
+impl fmt::Debug for MerkleHash {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "MerkleHash({self})")
+    }
+}
+
+impl fmt::Display for MerkleHash {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut buf = [0u8; 64];
+        let hex_string = base16ct::lower::encode_str(&self.0, &mut buf)
+            .expect("Buffer is sufficient for a SHA-256 hash");
+        write!(f, "{hex_string}")
+    }
+}
+
+#[cfg(test)]
+#[path = "hash_test.rs"]
+mod hash_test;

crates/ai/src/index/full_source_code_embedding/merkle_tree/hash_test.rs

@@ -0,0 +1,96 @@
+use super::MerkleHash;
+use crate::index::full_source_code_embedding::chunker::Fragment;
+use std::path::{Path, PathBuf};
+
+#[test]
+fn test_fragment_hash_from_content() {
+    let content = "fn main() { println!(\"Hello, world!\"); }";
+    let fragment = Fragment {
+        content,
+        file_path: Path::new("/foo/bar/bazz"),
+        start_line: 0,
+        end_line: 0,
+        start_byte_index: 0.into(),
+        end_byte_index: content.len().into(),
+    };
+    let hash = MerkleHash::from_fragment(&fragment);
+
+    assert_eq!(
+        "bb343b0950832ccd077f1515e842196f2ae4bb9e9261b0935ac57916c3cf305d",
+        hash.to_string()
+    );
+}
+
+#[test]
+fn test_node_hash_from_children() {
+    let path = PathBuf::from("/foo/bar/bazz");
+    let content1 = "fn func1() -> int { 1 }";
+    let content2 = "fn func2() -> int { 2 }";
+    let content3 = "fn func3() -> int { 3 }";
+
+    let leaf1 = MerkleHash::from_fragment(&Fragment {
+        content: content1,
+        file_path: path.as_path(),
+        start_line: 0,
+        end_line: 0,
+        start_byte_index: 0.into(),
+        end_byte_index: content1.len().into(),
+    });
+    let leaf2 = MerkleHash::from_fragment(&Fragment {
+        content: content2,
+        file_path: path.as_path(),
+        start_line: 1,
+        end_line: 1,
+        start_byte_index: content1.len().into(),
+        end_byte_index: (content1.len() + content2.len()).into(),
+    });
+    let leaf3 = MerkleHash::from_fragment(&Fragment {
+        content: content3,
+        file_path: path.as_path(),
+        start_line: 2,
+        end_line: 2,
+        start_byte_index: (content1.len() + content2.len()).into(),
+        end_byte_index: (content1.len() + content2.len() + content3.len()).into(),
+    });
+
+    // Create an iterator with the leaf hashes
+    let leaves = vec![&leaf1, &leaf2, &leaf3];
+    let hash = MerkleHash::from_hashes(leaves.into_iter());
+
+    assert_eq!(
+        "99c2f5b808870e4b1fcf163efbb588b6d5e074658d4ac43bab2eb5ffbe72c5cd",
+        hash.to_string()
+    );
+}
+
+#[test]
+fn test_merkle_hash_serialization_deserialization() {
+    // Create a MerkleHash from a fragment
+    let content = "fn main() { println!(\"Hello, world!\"); }";
+    let fragment = Fragment {
+        content,
+        file_path: Path::new("/foo/bar/bazz"),
+        start_line: 0,
+        end_line: 0,
+        start_byte_index: 0.into(),
+        end_byte_index: content.len().into(),
+    };
+    let original_hash = MerkleHash::from_fragment(&fragment);
+
+    // Serialize the hash to a JSON string
+    let serialized = serde_json::to_string(&original_hash).expect("Failed to serialize MerkleHash");
+
+    // Ensure serialized output is a hex string
+    assert!(serialized.starts_with("\""));
+    assert!(serialized.ends_with("\""));
+
+    // Deserialize back to a MerkleHash
+    let deserialized_hash: MerkleHash =
+        serde_json::from_str(&serialized).expect("Failed to deserialize MerkleHash");
+
+    // Verify deserialized hash matches the original
+    assert_eq!(original_hash, deserialized_hash);
+
+    // Also verify string representation matches
+    assert_eq!(original_hash.to_string(), deserialized_hash.to_string());
+}

crates/ai/src/index/full_source_code_embedding/merkle_tree/mod.rs

@@ -0,0 +1,32 @@
+use super::{chunker::Fragment, Error};
+
+mod hash;
+mod node;
+mod serialized_tree;
+mod tree;
+
+pub(super) use hash::MerkleHash;
+pub use hash::{ContentHash, NodeHash};
+pub(super) use node::NodeLens;
+pub(super) use serialized_tree::SerializedCodebaseIndex;
+pub(super) use tree::MerkleTree;
+
+use crate::index::Entry;
+cfg_if::cfg_if! {
+    if #[cfg(feature = "local_fs")] {
+        pub(super) use node::NodeId;
+        pub(super) use tree::TreeUpdateResult;
+    }
+}
+
+#[derive(Debug)]
+enum DirEntryOrFragment<'a> {
+    Entry(Entry),
+    Fragment(Fragment<'a>),
+}
+
+#[cfg(test)]
+mod test_util;
+
+#[cfg(test)]
+pub(super) use test_util::construct_test_merkle_tree;

crates/ai/src/index/full_source_code_embedding/merkle_tree/node.rs

@@ -0,0 +1,711 @@
+use crate::index::{
+    THREADPOOL, {DirectoryEntry, Entry, FileMetadata},
+};
+use anyhow::anyhow;
+
+use chrono::{DateTime, Utc};
+use itertools::Itertools;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use repo_metadata::entry::is_file_parsable;
+use serde::{Deserialize, Serialize};
+use sha2::{Digest, Sha256};
+use warp_util::standardized_path::StandardizedPath;
+
+use std::{
+    collections::{HashMap, HashSet},
+    ops::Range,
+    path::{Path, PathBuf},
+};
+use string_offset::ByteOffset;
+
+use crate::index::full_source_code_embedding::{
+    chunker::chunk_code,
+    fragment_metadata::{FragmentMetadata, LeafToFragmentMetadataUpdates},
+    Error,
+};
+
+use super::{
+    hash::MerkleHash,
+    serialized_tree::{SerializedFilesystemInfo, SerializedMerkleNode},
+    tree::UpdateFileResult,
+    ContentHash, DirEntryOrFragment, NodeHash,
+};
+
+/// ID that uniquely identifies a node in the merkle tree. It contains the node type
+/// as well as metadata that distinguishes nodes of the same type.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub(crate) enum NodeId {
+    /// A file node that contains fragment children
+    File {
+        absolute_path: PathBuf,
+        file_size: usize,
+        fs_modified_time: DateTime<Utc>,
+        file_contents_hash: String,
+    },
+    /// A directory node that contains file and directory children
+    Directory { absolute_path: PathBuf },
+    /// A leaf node representing a code fragment
+    Fragment {
+        absolute_path: PathBuf,
+        content_range: Range<ByteOffset>,
+    },
+}
+
+impl NodeId {
+    fn absolute_path(&self) -> &PathBuf {
+        match self {
+            Self::Directory { absolute_path } => absolute_path,
+            Self::File { absolute_path, .. } => absolute_path,
+            Self::Fragment { absolute_path, .. } => absolute_path,
+        }
+    }
+}
+
+/// A given node in the [`MerkleTree`].
+#[derive(Debug)]
+pub(super) struct MerkleNode {
+    /// The hash for the current node of the Merkle tree.
+    hash: MerkleHash,
+    /// The children of this merkle node.
+    children: Vec<MerkleNode>,
+    /// The ID of this node.
+    node_id: NodeId,
+}
+
+impl MerkleNode {
+    pub(super) fn new(
+        entry: DirEntryOrFragment<'_>,
+    ) -> Result<(MerkleNode, LeafToFragmentMetadataUpdates), Error> {
+        match entry {
+            DirEntryOrFragment::Entry(Entry::File(file)) => {
+                let local_path = file.path.to_local_path_lossy();
+                if !is_file_parsable(&local_path)? {
+                    return Err(Error::FileSizeExceeded);
+                }
+                let (file_size, fs_modified_time) = match std::fs::metadata(&local_path) {
+                    Ok(metadata) => {
+                        let file_size = metadata.len() as usize;
+                        if let Ok(fs_modified_time) = metadata.modified() {
+                            // Convert the SystemTime to DateTime<Utc>
+                            let fs_modified_time = fs_modified_time.into();
+
+                            (file_size, fs_modified_time)
+                        } else {
+                            log::warn!("Failed to get modified time for file {}", file.path);
+                            return Err(Error::FailedToGetMetadata(local_path));
+                        }
+                    }
+                    Err(_) => {
+                        log::warn!("Failed to get metadata for file {}", file.path);
+                        return Err(Error::FailedToGetMetadata(local_path));
+                    }
+                };
+
+                let file_contents = std::fs::read_to_string(&local_path)?;
+
+                // Compute SHA-256 hash of the file contents
+                let mut hasher = Sha256::new();
+                hasher.update(file_contents.as_bytes());
+                let file_contents_hash = format!("{:x}", hasher.finalize());
+
+                let fragments = chunk_code(&file_contents, &local_path);
+
+                let (children, mapping_updates): (Vec<_>, LeafToFragmentMetadataUpdates) =
+                    fragments
+                        .into_iter()
+                        .filter_map(|fragment| {
+                            Self::new(DirEntryOrFragment::Fragment(fragment)).ok()
+                        })
+                        .unzip();
+                if children.is_empty() {
+                    log::debug!(
+                        "Found empty file {} when generating the merkle tree",
+                        file.path
+                    );
+                    return Err(Error::EmptyNodeContent);
+                }
+
+                // Create a hash from all of the fragments of the file. We actively _do not_ sort here as the fragments
+                // are an ordered function of the file content.
+                let hash = MerkleHash::from_hashes(children.iter().map(|child| &child.hash));
+
+                // Add a small delay after file processing to reduce CPU spikes
+                // during large repository indexing, allowing other work to continue
+                std::thread::sleep(std::time::Duration::from_millis(5));
+
+                Ok((
+                    MerkleNode {
+                        hash,
+                        children,
+                        node_id: NodeId::File {
+                            absolute_path: file.path.to_local_path_lossy(),
+                            file_size,
+                            fs_modified_time,
+                            file_contents_hash,
+                        },
+                    },
+                    mapping_updates,
+                ))
+            }
+            DirEntryOrFragment::Entry(Entry::Directory(directory)) => {
+                let Some(pool) = THREADPOOL.as_ref() else {
+                    return Err(anyhow!("No threadpool exists for outline generation.").into());
+                };
+
+                let result = pool.install(|| {
+                    directory
+                        .children
+                        .into_par_iter()
+                        .filter_map(|node| Self::new(DirEntryOrFragment::Entry(node)).ok())
+                        .collect::<Vec<_>>()
+                });
+
+                let (mut children, mapping_updates): (Vec<_>, LeafToFragmentMetadataUpdates) =
+                    result.into_iter().unzip();
+
+                if children.is_empty() {
+                    return Err(Error::EmptyNodeContent);
+                }
+
+                // Sort the hashes to ensure we have consistent ordering for all files in the directory. We don't want a new
+                // hash if the `DirEntry`s are the same, but returned in a different order.
+                children.sort_unstable_by(|a, b| a.hash.cmp(&b.hash));
+
+                let hash = MerkleHash::from_hashes(children.iter().map(|child| &child.hash));
+                Ok((
+                    MerkleNode {
+                        hash,
+                        children,
+                        node_id: NodeId::Directory {
+                            absolute_path: directory.path.to_local_path_lossy(),
+                        },
+                    },
+                    mapping_updates,
+                ))
+            }
+            DirEntryOrFragment::Fragment(fragment) => {
+                if fragment.content.is_empty() {
+                    return Err(Error::EmptyNodeContent);
+                }
+                let hash = MerkleHash::from_fragment(&fragment);
+                let fragment_metadata = FragmentMetadata::from(&fragment);
+
+                let mut leaf_node_to_fragment_updates = LeafToFragmentMetadataUpdates::empty();
+                leaf_node_to_fragment_updates
+                    .to_insert
+                    .insert(hash.clone(), vec![fragment_metadata]);
+
+                Ok((
+                    MerkleNode {
+                        hash,
+                        children: vec![],
+                        node_id: NodeId::Fragment {
+                            absolute_path: fragment.file_path.to_path_buf(),
+                            content_range: fragment.start_byte_index..fragment.end_byte_index,
+                        },
+                    },
+                    leaf_node_to_fragment_updates,
+                ))
+            }
+        }
+    }
+
+    pub(super) fn from_serialized(
+        serialized_node: SerializedMerkleNode,
+        parent_path: &Path,
+    ) -> anyhow::Result<(MerkleNode, LeafToFragmentMetadataUpdates)> {
+        let hash = serialized_node.hash();
+
+        let mut children = vec![];
+        let mut leaf_node_to_fragment_updates = LeafToFragmentMetadataUpdates::empty();
+
+        let node_id = match serialized_node.fs_info {
+            SerializedFilesystemInfo::Directory { absolute_path } => {
+                NodeId::Directory { absolute_path }
+            }
+            SerializedFilesystemInfo::File {
+                absolute_path,
+                file_size,
+                fs_modified_time,
+                file_contents_hash,
+            } => NodeId::File {
+                absolute_path,
+                file_size,
+                fs_modified_time,
+                file_contents_hash,
+            },
+            SerializedFilesystemInfo::Fragment { location } => {
+                let file_path = parent_path.to_path_buf();
+                leaf_node_to_fragment_updates
+                    .to_insert
+                    .entry(hash.as_ref().clone())
+                    .or_default()
+                    .push(FragmentMetadata {
+                        absolute_path: file_path.clone(),
+                        location: (&location).into(),
+                    });
+
+                NodeId::Fragment {
+                    absolute_path: file_path.clone(),
+                    content_range: location.byte_range,
+                }
+            }
+        };
+
+        let absolute_path = node_id.absolute_path();
+
+        for child in serialized_node.children {
+            let (child_node, new_fragments) = Self::from_serialized(child, absolute_path)?;
+            leaf_node_to_fragment_updates.merge(new_fragments);
+            children.push(child_node);
+        }
+
+        Ok((
+            MerkleNode {
+                hash: hash.as_ref().clone(),
+                children,
+                node_id,
+            },
+            leaf_node_to_fragment_updates,
+        ))
+    }
+
+    // Recompute the hash for a given node. Note that for directories we expect the children to be sorted beforehand.
+    fn recompute_hash(&mut self) {
+        match &self.node_id {
+            NodeId::Directory { .. } | NodeId::File { .. } => {
+                self.hash = MerkleHash::from_hashes(self.children.iter().map(|child| &child.hash));
+            }
+            NodeId::Fragment { .. } => {
+                log::error!("Shouldn't need to recompute hash for fragments");
+            }
+        }
+    }
+
+    /// Remove a set of target file paths from this MerkleNode.
+    pub(super) fn remove_files(
+        &mut self,
+        paths: &mut HashSet<PathBuf>,
+        node_masks: &mut NodeMask,
+        node_to_fragment_updates: &mut LeafToFragmentMetadataUpdates,
+    ) -> UpdateFileResult {
+        match &self.node_id {
+            // Only visit a directory if it is the ancestor of the target path.
+            NodeId::Directory { absolute_path } => {
+                let mut paths_under_directory = filter_paths_under_directory(paths, absolute_path);
+
+                if paths_under_directory.is_empty() {
+                    return UpdateFileResult::NoChange;
+                }
+
+                // Track the indices that need to be removed.
+                let mut removal_idx = vec![];
+                let mut updated_idx = HashMap::new();
+                for (i, child) in self.children.iter_mut().enumerate() {
+                    let mut node = NodeMask::new(i);
+                    match child.remove_files(
+                        &mut paths_under_directory,
+                        &mut node,
+                        node_to_fragment_updates,
+                    ) {
+                        // If the node should be removed, remove it from the children.
+                        UpdateFileResult::Deleted => {
+                            removal_idx.push(i);
+                        }
+                        // If the node is updated, push the current index to node paths.
+                        UpdateFileResult::Updated => {
+                            updated_idx.insert(i - removal_idx.len(), node);
+                        }
+                        UpdateFileResult::NoChange => (),
+                    };
+
+                    if paths_under_directory.is_empty() {
+                        break;
+                    }
+                }
+
+                for idx in removal_idx.into_iter().rev() {
+                    self.children.remove(idx);
+                }
+
+                // If there is no more children to this merkle node, delete the current node.
+                if self.children.is_empty() {
+                    return UpdateFileResult::Deleted;
+                }
+
+                // We need to rebuild the children since the order of the children must be strictly sorted by their
+                // merkle hash.
+                let new_children_with_idx = self
+                    .children
+                    .drain(..)
+                    .enumerate()
+                    .sorted_by(|(_, a), (_, b)| a.hash.cmp(&b.hash));
+
+                for (new_idx, (old_idx, child)) in new_children_with_idx.enumerate() {
+                    if let Some(mut node) = updated_idx.remove(&old_idx) {
+                        node.index = new_idx;
+                        node_masks.add_child(node);
+                    }
+                    self.children.push(child);
+                }
+
+                if !updated_idx.is_empty() {
+                    log::error!("Updated index should be empty after an upsert request!");
+                }
+
+                // Otherwise, update the cache.
+                self.recompute_hash();
+                UpdateFileResult::Updated
+            }
+            // Only visit a file if it matches the target path.
+            NodeId::File { absolute_path, .. } if paths.remove(absolute_path) => {
+                node_to_fragment_updates.to_remove.insert(
+                    absolute_path.to_path_buf(),
+                    self.child_hashes().cloned().collect_vec(),
+                );
+                UpdateFileResult::Deleted
+            }
+            _ => UpdateFileResult::NoChange,
+        }
+    }
+
+    /// Update / insert a batch target file paths to this MerkleNode. Return whether the current node is updated.
+    pub(super) fn upsert_files(
+        &mut self,
+        paths: &mut HashSet<PathBuf>,
+        node_masks: &mut NodeMask,
+        leaf_node_to_fragment_updates: &mut LeafToFragmentMetadataUpdates,
+    ) -> UpdateFileResult {
+        match &self.node_id {
+            NodeId::Directory { absolute_path } => {
+                let mut paths_under_directory = filter_paths_under_directory(paths, absolute_path);
+
+                if paths_under_directory.is_empty() {
+                    return UpdateFileResult::NoChange;
+                }
+
+                let mut updated_idx = HashMap::new();
+                let mut removal_idx = vec![];
+
+                for (i, child) in self.children.iter_mut().enumerate() {
+                    let mut node = NodeMask::new(i);
+                    let updated = child.upsert_files(
+                        &mut paths_under_directory,
+                        &mut node,
+                        leaf_node_to_fragment_updates,
+                    );
+
+                    // Only update node_masks if the child is updated.
+                    match updated {
+                        UpdateFileResult::Deleted => removal_idx.push(i),
+                        UpdateFileResult::Updated => {
+                            updated_idx.insert(i - removal_idx.len(), node);
+                        }
+                        UpdateFileResult::NoChange => (),
+                    }
+
+                    // There is no more things to update. We could break early.
+                    if paths_under_directory.is_empty() {
+                        break;
+                    }
+                }
+
+                for idx in removal_idx.into_iter().rev() {
+                    self.children.remove(idx);
+                }
+
+                // If there is no more children to this merkle node and we are not creating new nodes, delete the current node.
+                if self.children.is_empty() && paths_under_directory.is_empty() {
+                    return UpdateFileResult::Deleted;
+                }
+
+                let mut entries_to_create = Vec::new();
+
+                // If none of the existing child could be updated these, we need to create new nodes.
+                // The algorithm works as below:
+                // 1) Convert the to-be-inserted paths into a chain of ancestors (e.g. a/b/c -> [a, a/b, a/b/c]).
+                // 2) Dedupe the ancestors using a mapping of path -> directory entry.
+                // 3) Attach each directory node to its parent node.
+                // 4) Add the root-level directory nodes to entries_to_create.
+                //
+                // Note that for file nodes, we add them directly to entries_to_create if it is on the root level and
+                // to its corresponding parent node otherwise.
+                let mut created_dirs: HashMap<PathBuf, DirectoryEntry> = HashMap::new();
+
+                // First pass: Create all directory entries with empty children lists
+                for path in paths_under_directory {
+                    // Skip upserting full or non-existent directories.
+                    if path.is_dir() || !path.exists() {
+                        continue;
+                    }
+
+                    // Get all parent directories that need to be created
+                    let mut ancestors = Vec::new();
+                    let mut current = path.clone();
+
+                    // Start from the path's parent and work up to but not including absolute_path
+                    while let Some(parent) = current.parent() {
+                        current = parent.to_path_buf();
+
+                        if current == *absolute_path {
+                            break;
+                        }
+
+                        ancestors.push(current.clone());
+                    }
+
+                    if current != *absolute_path {
+                        log::warn!("Path should match absolute path before None");
+                        continue;
+                    }
+
+                    // Process ancestors from deepest to shallowest (reverse order)
+                    ancestors.reverse();
+
+                    // Create intermediate directories if they don't exist yet
+                    for ancestor in ancestors {
+                        created_dirs
+                            .entry(ancestor.clone())
+                            .or_insert_with(|| DirectoryEntry {
+                                path: StandardizedPath::try_from_local(&ancestor)
+                                    .expect("ancestor paths are always absolute"),
+                                children: Vec::new(),
+                                ignored: false,
+                                loaded: false,
+                            });
+                    }
+
+                    // Add the file entry to its parent directory's children list
+                    if let Some(parent_path) = path.parent() {
+                        if parent_path != absolute_path {
+                            if let Some(parent_dir) = created_dirs.get_mut(parent_path) {
+                                parent_dir
+                                    .children
+                                    .push(Entry::File(FileMetadata::new(path, false)));
+                            }
+                        } else {
+                            entries_to_create.push(Entry::File(FileMetadata::new(path, false)));
+                        }
+                    }
+                }
+
+                // Second pass: Sort directories by depth (deepest first) and establish relationships from bottom up
+                let mut dir_paths: Vec<PathBuf> = created_dirs.keys().cloned().collect();
+                // Sort by component count in reverse order - deeper paths first
+                dir_paths.sort_by_key(|p| std::cmp::Reverse(p.components().count()));
+
+                for dir_path in dir_paths {
+                    if let Some(parent_path) = dir_path.parent() {
+                        if let Some(child_dir) = created_dirs.remove(&dir_path) {
+                            // Skip if parent is the root directory
+                            if parent_path != absolute_path {
+                                if let Some(parent_dir) = created_dirs.get_mut(parent_path) {
+                                    // Now we have the completed child directory with all its children
+                                    parent_dir.children.push(Entry::Directory(child_dir));
+                                }
+                            } else {
+                                entries_to_create.push(Entry::Directory(child_dir));
+                            }
+                        }
+                    }
+                }
+
+                for entry in entries_to_create {
+                    let res = MerkleNode::new(DirEntryOrFragment::Entry(entry));
+                    let (child, mapping_updates) = match res {
+                        Ok((child, mapping)) => (child, mapping),
+                        // When encountering a node construction error, instead of early returning and interrupting the rest of the update,
+                        // consider it a skippable error.
+                        // TODO: We should capture and log these errors in the telemetry.
+                        Err(e) => {
+                            log::debug!("Failed to create new node for update: {e:#}");
+                            continue;
+                        }
+                    };
+
+                    self.children.push(child);
+                    leaf_node_to_fragment_updates.merge(mapping_updates);
+                    updated_idx.insert(self.children.len() - 1, NodeMask::new(0));
+                }
+
+                // We need to rebuild the children since the order of the children must be strictly sorted by their
+                // merkle hash.
+                let new_children_with_idx = self
+                    .children
+                    .drain(..)
+                    .enumerate()
+                    .sorted_by(|(_, a), (_, b)| a.hash.cmp(&b.hash));
+
+                for (new_idx, (old_idx, child)) in new_children_with_idx.enumerate() {
+                    if let Some(mut node) = updated_idx.remove(&old_idx) {
+                        node.index = new_idx;
+                        node_masks.add_child(node);
+                    }
+                    self.children.push(child);
+                }
+
+                if !updated_idx.is_empty() {
+                    log::error!("Updated index should be empty after an upsert request!");
+                }
+                self.recompute_hash();
+                UpdateFileResult::Updated
+            }
+            // For files, only a single path can match at a single time.
+            NodeId::File { absolute_path, .. } if paths.remove(absolute_path) => {
+                leaf_node_to_fragment_updates.to_remove.insert(
+                    absolute_path.clone(),
+                    self.child_hashes().cloned().collect_vec(),
+                );
+
+                if !absolute_path.exists() {
+                    return UpdateFileResult::Deleted;
+                }
+
+                let (new_node, mapping_update) = match MerkleNode::new(DirEntryOrFragment::Entry(
+                    Entry::File(FileMetadata::new(absolute_path.clone(), false)),
+                )) {
+                    Ok(res) => res,
+                    // If we run into a file permission error / empty node / exceeded max file limit, delete the node since we can't
+                    // determine what's the updated content.
+                    Err(_) => return UpdateFileResult::Deleted,
+                };
+                leaf_node_to_fragment_updates.merge(mapping_update);
+                self.children = new_node.children;
+                self.hash = new_node.hash;
+                self.node_id = new_node.node_id;
+                UpdateFileResult::Updated
+            }
+            _ => UpdateFileResult::NoChange,
+        }
+    }
+
+    pub(super) fn absolute_path(&self) -> &Path {
+        self.node_id.absolute_path()
+    }
+
+    fn child_hashes(&self) -> impl Iterator<Item = &MerkleHash> {
+        self.children.iter().map(|child| &child.hash)
+    }
+
+    pub(super) fn children(&self) -> impl Iterator<Item = &MerkleNode> {
+        self.children.iter()
+    }
+
+    pub(super) fn count_children(&self) -> usize {
+        self.children.len()
+    }
+
+    pub(super) fn child_at(&self, index: usize) -> &MerkleNode {
+        &self.children[index]
+    }
+
+    pub(super) fn hash(&self) -> &MerkleHash {
+        &self.hash
+    }
+
+    pub(super) fn node_id(&self) -> &NodeId {
+        &self.node_id
+    }
+
+    pub(super) fn is_fragment(&self) -> bool {
+        matches!(self.node_id, NodeId::Fragment { .. })
+    }
+}
+
+fn filter_paths_under_directory(
+    paths: &mut HashSet<PathBuf>,
+    curr_path: &PathBuf,
+) -> HashSet<PathBuf> {
+    let mut paths_under_directory = HashSet::new();
+
+    // Construct and filter out paths that are under the current directory.
+    for path in paths.iter() {
+        if path.starts_with(curr_path) {
+            paths_under_directory.insert(path.clone());
+        }
+    }
+
+    for filter_path in paths_under_directory.iter() {
+        paths.remove(filter_path);
+    }
+
+    paths_under_directory
+}
+
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct NodeLens<'a> {
+    node: &'a MerkleNode,
+}
+
+impl<'a> NodeLens<'a> {
+    pub(super) fn new(node: &'a MerkleNode) -> Self {
+        Self { node }
+    }
+
+    pub fn children(&self) -> impl Iterator<Item = NodeLens<'a>> {
+        self.node.children().map(|node| NodeLens { node })
+    }
+
+    pub fn hash(&self) -> NodeHash {
+        NodeHash::new(self.node.hash().clone())
+    }
+
+    pub fn content_hash(&self) -> Option<ContentHash> {
+        self.is_leaf()
+            .then(|| ContentHash::new(self.node.hash().clone()))
+    }
+
+    pub fn is_leaf(&self) -> bool {
+        self.node.is_fragment()
+    }
+
+    pub fn path(&self) -> &Path {
+        self.node.absolute_path()
+    }
+
+    pub(crate) fn node_id(&self) -> &NodeId {
+        self.node.node_id()
+    }
+}
+
+#[derive(Default)]
+pub(super) enum ChildrenPath {
+    #[default]
+    All,
+    SpecificChildren(Vec<NodeMask>),
+}
+
+/// A path from the leaf changed file node to the root.
+///
+/// Each NodeMask corresponds to a node in the tree. It contains the index of the node
+/// it is referencing in the parent node's children.
+///
+/// Note a NodeMask strictly couples with a snapshot of a Merkle tree. If the tree
+/// has been edited afterwards, the NodeMask will no longer be valid.
+#[derive(Default)]
+pub(super) struct NodeMask {
+    pub(super) index: usize,
+    pub(super) children: ChildrenPath,
+}
+
+impl NodeMask {
+    pub(super) fn new(index: usize) -> Self {
+        Self {
+            index,
+            children: Default::default(),
+        }
+    }
+
+    pub(super) fn add_child(&mut self, node: Self) {
+        match &mut self.children {
+            ChildrenPath::All => {
+                self.children = ChildrenPath::SpecificChildren(vec![node]);
+            }
+            ChildrenPath::SpecificChildren(children) => children.push(node),
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "node_test.rs"]
+mod tests;

crates/ai/src/index/full_source_code_embedding/merkle_tree/node_test.rs

@@ -0,0 +1,835 @@
+use crate::index::full_source_code_embedding::{
+    fragment_metadata::LeafToFragmentMetadataUpdates, merkle_tree::DirEntryOrFragment,
+};
+use repo_metadata::{DirectoryEntry, Entry};
+use virtual_fs::{Stub, VirtualFS};
+
+use std::collections::HashSet;
+
+use super::{MerkleNode, NodeMask};
+
+/// Tests that node hashes for directories are sorted (meaning they are resilient to files within
+/// the directory being in a different order).
+#[test]
+fn test_node_hash_for_directory_is_sorted() {
+    VirtualFS::test(
+        "test_node_hash_for_directory_is_sorted",
+        |dirs, mut sandbox| {
+            sandbox.with_files(vec![Stub::FileWithContent("foo", "foo")]);
+            sandbox.with_files(vec![Stub::FileWithContent("bar", "bar")]);
+
+            let mut directory_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                    .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: false,
+            };
+
+            for file in ["foo", "bar"] {
+                directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            let (node, _leaf_to_fragment_updates) = MerkleNode::new(DirEntryOrFragment::Entry(
+                Entry::Directory(directory_entry.clone()),
+            ))
+            .expect("Should be able to construct node");
+
+            directory_entry.children.clear();
+
+            for file in ["bar", "foo"] {
+                directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            let (node_reverse, _leaf_to_fragment_updates) =
+                MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(directory_entry)))
+                    .expect("Should be able to construct node");
+
+            // The node hashes should be the same even though the files were returned in different orders.
+            assert_eq!(node.hash, node_reverse.hash);
+        },
+    );
+}
+
+/// Tests that upserting a file updates the Merkle tree correctly.
+#[test]
+fn test_merkle_node_upsert_file() {
+    VirtualFS::test("test_merkle_node_upsert_file", |dirs, mut sandbox| {
+        // Create a directory with an initial file
+        sandbox.with_files(vec![Stub::FileWithContent(
+            "initial.txt",
+            "initial content",
+        )]);
+
+        let mut directory_entry = DirectoryEntry {
+            path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                .unwrap(),
+            children: vec![],
+            ignored: false,
+            loaded: true,
+        };
+
+        // Add the initial file to the directory entry
+        directory_entry
+            .find_or_insert_child(&dirs.tests().join("initial.txt"))
+            .expect("Should be able to insert into directory entry");
+
+        // Create the initial MerkleNode
+        let (mut node, initial_metadata_update) = MerkleNode::new(DirEntryOrFragment::Entry(
+            Entry::Directory(directory_entry.clone()),
+        ))
+        .expect("Should be able to construct node");
+
+        assert_eq!(
+            initial_metadata_update.to_insert.len(),
+            1,
+            "Should insert one file's metadata",
+        );
+        assert!(
+            initial_metadata_update.to_remove.is_empty(),
+            "Should not remove any file metadata",
+        );
+
+        let initial_root_hash = node.hash.clone();
+
+        // Create a new file to upsert
+        sandbox.with_files(vec![Stub::FileWithContent("new.txt", "new content")]);
+        directory_entry
+            .find_or_insert_child(&dirs.tests().join("new.txt"))
+            .expect("Should be able to insert into directory entry");
+
+        // Upsert the new file
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.upsert_files(
+            &mut HashSet::from([dirs.tests().join("new.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        // Verify the hash changed after adding a new file
+        assert_ne!(
+            initial_root_hash, node.hash,
+            "Hash should change after adding a new file",
+        );
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_insert.len(),
+            1,
+            "Upserting a new file should insert its content into metadata mapping",
+        );
+        assert!(
+            leaf_to_fragment_metadata_updates.to_remove.is_empty(),
+            "Inserting a new file should not remove any content from metadata mapping",
+        );
+
+        // Updated hash should be the same as reconstructing hash from scratch.
+        let hash_from_scratch = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+            directory_entry.clone(),
+        )))
+        .expect("Should be able to construct node")
+        .0
+        .hash;
+        assert_eq!(node.hash, hash_from_scratch);
+
+        // Remember the hash after adding the new file
+        let hash_after_add = node.hash.clone();
+
+        // Modify an existing file
+        sandbox.with_files(vec![Stub::FileWithContent(
+            "initial.txt",
+            "modified content",
+        )]);
+
+        // Upsert the modified file
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.upsert_files(
+            &mut HashSet::from([dirs.tests().join("initial.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        let updated_leaf_hash = leaf_to_fragment_metadata_updates
+            .to_insert
+            .keys()
+            .next()
+            .unwrap();
+        let updated_metadata_entry =
+            leaf_to_fragment_metadata_updates.to_insert[updated_leaf_hash].clone();
+
+        // Verify the hash changed after modifying a file
+        assert_ne!(
+            hash_after_add, node.hash,
+            "Hash should change after modifying a file"
+        );
+
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_insert.len(),
+            1,
+            "Upserting a modified file should insert its content into metadata mapping",
+        );
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_remove.len(),
+            1,
+            "Upserting a modified file should remove its old content from metadata mapping",
+        );
+
+        // Updated hash should be the same as reconstructing hash from scratch.
+        let hash_from_scratch = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+            directory_entry.clone(),
+        )))
+        .expect("Should be able to construct node")
+        .0
+        .hash;
+        assert_eq!(node.hash, hash_from_scratch);
+
+        // Remember the hash after modification
+        let hash_after_modify = node.hash.clone();
+
+        // Upsert with the same content (should not change the hash)
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.upsert_files(
+            &mut HashSet::from([dirs.tests().join("initial.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        // Verify the hash did not change when content is the same
+        assert_eq!(
+            hash_after_modify, node.hash,
+            "Hash should not change when upserting the same content"
+        );
+
+        // Verify that we remove and re-insert the file metadata in the leaf mapping.
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_remove.len(),
+            1,
+            "Should remove an entry from the metadata mapping",
+        );
+        assert!(
+            leaf_to_fragment_metadata_updates
+                .to_remove
+                .contains_key(&dirs.tests().join("initial.txt")),
+            "Upserting the same content should remove the old entry from the metadata mapping",
+        );
+
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_insert.len(),
+            1,
+            "Should re-insert an entry into the metadata mapping",
+        );
+        assert!(
+            leaf_to_fragment_metadata_updates
+                .to_insert
+                .contains_key(updated_leaf_hash),
+            "Upserting the same content re-inserts the same hash into the metadata mapping",
+        );
+        assert_eq!(
+            leaf_to_fragment_metadata_updates
+                .to_insert
+                .get(updated_leaf_hash)
+                .unwrap(),
+            &updated_metadata_entry,
+            "Upserting the same content re-inserts the same metadata into the metadata mapping"
+        );
+
+        // Test upserting a file in a new subdirectory that doesn't exist yet
+        let hash_before_subdirectory = node.hash.clone();
+
+        // Create a new subdirectory structure with a file
+        sandbox.mkdir("a");
+        sandbox.with_files(vec![Stub::FileWithContent(
+            "a/b.txt",
+            "subdirectory file content",
+        )]);
+
+        // Upsert the file in the new subdirectory
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.upsert_files(
+            &mut HashSet::from([dirs.tests().join("a/b.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        // Verify the hash changed after adding a file in a new subdirectory
+        assert_ne!(
+            hash_before_subdirectory, node.hash,
+            "Hash should change after adding a file in a new subdirectory"
+        );
+
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_insert.len(),
+            1,
+            "Upserting a file in a new subdirectory should insert its content into metadata mapping"
+        );
+        assert!(
+            leaf_to_fragment_metadata_updates.to_remove.is_empty(),
+            "Upserting a file in a new subdirectory should not remove any content from metadata mapping"
+        );
+
+        // Create a new MerkleTree with the expected structure manually
+        let mut expected_directory_entry = DirectoryEntry {
+            path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                .unwrap(),
+            children: vec![],
+            ignored: false,
+            loaded: true,
+        };
+
+        // Add both the existing files and new subdirectory structure
+        for file in ["initial.txt", "new.txt"] {
+            expected_directory_entry
+                .find_or_insert_child(&dirs.tests().join(file))
+                .expect("Should be able to insert into directory entry");
+        }
+
+        // Create a subdirectory entry for 'a'
+        let mut subdirectory_entry = DirectoryEntry {
+            path: warp_util::standardized_path::StandardizedPath::try_from_local(
+                &dirs.tests().join("a"),
+            )
+            .unwrap(),
+            children: vec![],
+            ignored: false,
+            loaded: true,
+        };
+
+        // Add the 'b.txt' file to the subdirectory
+        subdirectory_entry
+            .find_or_insert_child(&dirs.tests().join("a/b.txt"))
+            .expect("Should be able to insert into subdirectory entry");
+
+        // Add the subdirectory to the main directory
+        expected_directory_entry
+            .children
+            .push(Entry::Directory(subdirectory_entry));
+
+        // Create a MerkleNode with the expected structure
+        let (expected_node, _) = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+            expected_directory_entry,
+        )))
+        .expect("Should be able to construct node with expected structure");
+
+        // The hash of our modified node should match the hash of the manually constructed node
+        assert_eq!(
+            node.hash, expected_node.hash,
+            "Hash of node with upserted subdirectory should match hash of node constructed with expected structure"
+        );
+    });
+}
+
+/// Tests that removing a file updates the Merkle tree correctly.
+#[test]
+fn test_merkle_node_remove_file() {
+    VirtualFS::test("test_merkle_node_remove_file", |dirs, mut sandbox| {
+        // Create a directory with multiple files
+        sandbox.with_files(vec![
+            Stub::FileWithContent("file1.txt", "content 1"),
+            Stub::FileWithContent("file2.txt", "content 2"),
+            Stub::FileWithContent("file3.txt", "content 3"),
+        ]);
+
+        let mut directory_entry = DirectoryEntry {
+            path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                .unwrap(),
+            children: vec![],
+            ignored: false,
+            loaded: true,
+        };
+
+        // Add all files to the directory entry
+        for file in ["file1.txt", "file2.txt", "file3.txt"] {
+            directory_entry
+                .find_or_insert_child(&dirs.tests().join(file))
+                .expect("Should be able to insert into directory entry");
+        }
+
+        // Create the initial MerkleNode
+        let (mut node, initial_metadata_updates) = MerkleNode::new(DirEntryOrFragment::Entry(
+            Entry::Directory(directory_entry.clone()),
+        ))
+        .expect("Should be able to construct node");
+
+        assert_eq!(
+            initial_metadata_updates.to_insert.len(),
+            3,
+            "Should insert three files' metadata"
+        );
+        assert!(
+            initial_metadata_updates.to_remove.is_empty(),
+            "Should not remove any file metadata"
+        );
+
+        let initial_hash = node.hash.clone();
+
+        // Remove one of the files
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.remove_files(
+            &mut HashSet::from([dirs.tests().join("file2.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        // Verify the hash changed after removing a file
+        assert_ne!(
+            initial_hash, node.hash,
+            "Hash should change after removing a file"
+        );
+
+        // Verify that we update the fragment metadata
+        assert_eq!(
+            leaf_to_fragment_metadata_updates.to_remove.len(),
+            1,
+            "Removing a file should remove its content from metadata mapping"
+        );
+        assert!(
+            leaf_to_fragment_metadata_updates.to_insert.is_empty(),
+            "Removing a file should not insert any new metadata"
+        );
+
+        // Remember the hash after removal
+        let hash_after_remove = node.hash.clone();
+
+        // Try removing a non-existent file (should not change the hash)
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+        node.remove_files(
+            &mut HashSet::from([dirs.tests().join("nonexistent.txt")]),
+            &mut node_path,
+            &mut leaf_to_fragment_metadata_updates,
+        );
+
+        // Verify the hash did not change when removing a non-existent file
+        assert_eq!(
+            hash_after_remove, node.hash,
+            "Hash should not change when removing a non-existent file"
+        );
+
+        // Verify that there are no metadata updates
+        assert!(
+            leaf_to_fragment_metadata_updates.is_empty(),
+            "Metadata updates should be empty after trying to remove a non-existent file"
+        );
+
+        // Create a new MerkleNode with only the remaining files
+        let mut directory_entry_after_remove = DirectoryEntry {
+            path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                .unwrap(),
+            children: vec![],
+            ignored: false,
+            loaded: true,
+        };
+
+        for file in ["file1.txt", "file3.txt"] {
+            directory_entry_after_remove
+                .find_or_insert_child(&dirs.tests().join(file))
+                .expect("Should be able to insert into directory entry");
+        }
+
+        let (node_after_remove, _) = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+            directory_entry_after_remove,
+        )))
+        .expect("Should be able to construct node with remaining files");
+
+        // Verify that manually constructing a node without the removed file
+        // produces the same hash as removing the file from an existing node
+        assert_eq!(
+            node.hash, node_after_remove.hash,
+            "Hash after remove should match hash of node constructed without the file"
+        );
+    });
+}
+
+/// Tests that upserting and removing multiple files updates the Merkle tree correctly.
+#[test]
+fn test_merkle_node_multiple_operations() {
+    VirtualFS::test(
+        "test_merkle_node_multiple_operations",
+        |dirs, mut sandbox| {
+            // Create a directory with multiple initial files
+            sandbox.with_files(vec![
+                Stub::FileWithContent("file1.txt", "content 1"),
+                Stub::FileWithContent("file2.txt", "content 2"),
+                Stub::FileWithContent("file3.txt", "content 3"),
+            ]);
+
+            let mut directory_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                    .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: true,
+            };
+
+            // Add all initial files to the directory entry
+            for file in ["file1.txt", "file2.txt", "file3.txt"] {
+                directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            // Create the initial MerkleNode
+            let (mut node, initial_metadata_updates) = MerkleNode::new(DirEntryOrFragment::Entry(
+                Entry::Directory(directory_entry.clone()),
+            ))
+            .expect("Should be able to construct node");
+
+            // Ensure the initial leaf-node-to-fragment-metadata updates are correct.
+            assert_eq!(
+                initial_metadata_updates.to_insert.len(),
+                3,
+                "Should insert three files' metadata"
+            );
+
+            let initial_hash = node.hash.clone();
+
+            // Test 1: Upsert multiple files at once
+            // Create multiple new files to upsert
+            sandbox.with_files(vec![
+                Stub::FileWithContent("file4.txt", "content 4"),
+                Stub::FileWithContent("file5.txt", "content 5"),
+            ]);
+            for file in ["file4.txt", "file5.txt"] {
+                directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            // Upsert multiple files at once
+            let mut node_path = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.upsert_files(
+                &mut HashSet::from([
+                    dirs.tests().join("file4.txt"),
+                    dirs.tests().join("file5.txt"),
+                ]),
+                &mut node_path,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Verify the hash changed after adding multiple files
+            assert_ne!(
+                initial_hash, node.hash,
+                "Hash should change after adding multiple files"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_insert.len(),
+                2,
+                "Upserting new files should insert their content into metadata mapping"
+            );
+            assert!(
+                leaf_to_fragment_metadata_updates.to_remove.is_empty(),
+                "Upserting new files should not remove content from metadata mapping"
+            );
+
+            // Updated hash should be the same as reconstructing hash from scratch.
+            let hash_from_scratch = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+                directory_entry.clone(),
+            )))
+            .expect("Should be able to construct node")
+            .0
+            .hash;
+            assert_eq!(node.hash, hash_from_scratch);
+
+            let hash_after_multiple_add = node.hash.clone();
+
+            // Test 2: Remove multiple files at once
+            let mut node_path = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.remove_files(
+                &mut HashSet::from([
+                    dirs.tests().join("file1.txt"),
+                    dirs.tests().join("file3.txt"),
+                ]),
+                &mut node_path,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Verify the hash changed after removing multiple files
+            assert_ne!(
+                hash_after_multiple_add, node.hash,
+                "Hash should change after removing multiple files"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_remove.len(),
+                2,
+                "Removing multiple files should remove their content from metadata mapping"
+            );
+            assert!(
+                leaf_to_fragment_metadata_updates.to_insert.is_empty(),
+                "Removing multiple files should not insert any new metadata"
+            );
+
+            let mut directory_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                    .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: true,
+            };
+            for file in ["file2.txt", "file4.txt", "file5.txt"] {
+                directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            // Updated hash should be the same as reconstructing hash from scratch.
+            let hash_from_scratch = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+                directory_entry.clone(),
+            )))
+            .expect("Should be able to construct node")
+            .0
+            .hash;
+            assert_eq!(node.hash, hash_from_scratch);
+
+            let hash_after_multiple_remove = node.hash.clone();
+
+            // Test 3: Mixed operations - modify an existing file and add a new file
+            // Modify an existing file
+            sandbox.with_files(vec![
+                Stub::FileWithContent("file2.txt", "modified content 2"),
+                Stub::FileWithContent("file6.txt", "content 6"),
+            ]);
+            directory_entry
+                .find_or_insert_child(&dirs.tests().join("file6.txt"))
+                .expect("Should be able to insert into directory entry");
+
+            let mut node_path = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.upsert_files(
+                &mut HashSet::from([
+                    dirs.tests().join("file2.txt"),
+                    dirs.tests().join("file6.txt"),
+                ]),
+                &mut node_path,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Verify the hash changed after mixed operations
+            assert_ne!(
+                hash_after_multiple_remove, node.hash,
+                "Hash should change after mixed upsert operations"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_insert.len(),
+                2,
+                "Upserting modified and new files should insert their content into metadata mapping"
+            );
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_remove.len(),
+                1,
+                "Upserting modified and new files should remove the old content of the modified file from metadata mapping"
+            );
+
+            // Updated hash should be the same as reconstructing hash from scratch.
+            let hash_from_scratch = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+                directory_entry.clone(),
+            )))
+            .expect("Should be able to construct node")
+            .0
+            .hash;
+            assert_eq!(node.hash, hash_from_scratch);
+
+            let hash_after_mixed_upsert = node.hash.clone();
+
+            // Test 4: Edge case - upsert a file that already exists with the same content
+            let mut node_path = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.upsert_files(
+                &mut HashSet::from([dirs.tests().join("file2.txt")]),
+                &mut node_path,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Verify hash didn't change when upserting with the same content
+            assert_eq!(
+                hash_after_mixed_upsert, node.hash,
+                "Hash should not change when upserting files with the same content"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_remove.len(),
+                1,
+                "Upserting files with the same content first removes it from the metadata mapping"
+            );
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_insert.len(),
+                1,
+                "Upserting files with the same content re-adds it to the metadata mapping"
+            );
+
+            // Test 5: Edge case - remove a non-existent file while upserting a new file
+            // Create nested directories and a new file
+            sandbox.mkdir("subdir1");
+            sandbox.mkdir("subdir1/subdir2");
+            sandbox.with_files(vec![Stub::FileWithContent(
+                "subdir1/subdir2/nested.txt",
+                "nested content",
+            )]);
+
+            // Do mixed operations - remove non-existent file while upserting a new file in a subdirectory
+            let mut node_path_remove = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.remove_files(
+                &mut HashSet::from([dirs.tests().join("nonexistent.txt")]),
+                &mut node_path_remove,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Hash shouldn't change after trying to remove a non-existent file
+            assert_eq!(
+                hash_after_mixed_upsert, node.hash,
+                "Hash should not change when removing a non-existent file"
+            );
+
+            assert!(
+                leaf_to_fragment_metadata_updates.is_empty(),
+                "Removing a non-existent file should not modify metadata mapping"
+            );
+
+            // Add the nested file
+            let mut node_path_upsert = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.upsert_files(
+                &mut HashSet::from([dirs.tests().join("subdir1/subdir2/nested.txt")]),
+                &mut node_path_upsert,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Hash should change after adding the nested file
+            assert_ne!(
+                hash_after_mixed_upsert, node.hash,
+                "Hash should change after adding a file in a nested subdirectory"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_insert.len(),
+                1,
+                "Upserting a file in a nested subdirectory should insert its content into metadata mapping",
+            );
+            assert!(
+                leaf_to_fragment_metadata_updates.to_remove.is_empty(),
+                "Upserting a file in a nested subdirectory should not remove any files from metadata mapping",
+            );
+
+            let hash_after_nested_add = node.hash.clone();
+
+            // Test 6: Mixed operations - remove multiple files and add multiple files at once
+            sandbox.with_files(vec![
+                Stub::FileWithContent("new_file1.txt", "new content 1"),
+                Stub::FileWithContent("new_file2.txt", "new content 2"),
+            ]);
+
+            // Remove some files
+            let mut node_path_remove = NodeMask::default();
+            let mut leaf_to_fragment_metadata_updates = LeafToFragmentMetadataUpdates::empty();
+            node.remove_files(
+                &mut HashSet::from([
+                    dirs.tests().join("file4.txt"),
+                    dirs.tests().join("file5.txt"),
+                ]),
+                &mut node_path_remove,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Add new files
+            let mut node_path_upsert = NodeMask::default();
+            node.upsert_files(
+                &mut HashSet::from([
+                    dirs.tests().join("new_file1.txt"),
+                    dirs.tests().join("new_file2.txt"),
+                ]),
+                &mut node_path_upsert,
+                &mut leaf_to_fragment_metadata_updates,
+            );
+
+            // Verify the hash changed after these mixed operations
+            assert_ne!(
+                hash_after_nested_add, node.hash,
+                "Hash should change after removing and adding multiple files"
+            );
+
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_insert.len(),
+                2,
+                "Upserting multiple files should insert their content into metadata mapping"
+            );
+            assert_eq!(
+                leaf_to_fragment_metadata_updates.to_remove.len(),
+                2,
+                "Removing multiple files should remove their content from metadata mapping"
+            );
+
+            // Create a new MerkleTree with the expected final structure manually
+            let mut expected_directory_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests())
+                    .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: true,
+            };
+
+            // Add all files that should be in the final structure
+            for file in ["file2.txt", "file6.txt", "new_file1.txt", "new_file2.txt"] {
+                expected_directory_entry
+                    .find_or_insert_child(&dirs.tests().join(file))
+                    .expect("Should be able to insert into directory entry");
+            }
+
+            // Create nested subdirectory structure
+            let mut subdir1_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(
+                    &dirs.tests().join("subdir1"),
+                )
+                .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: true,
+            };
+
+            let mut subdir2_entry = DirectoryEntry {
+                path: warp_util::standardized_path::StandardizedPath::try_from_local(
+                    &dirs.tests().join("subdir1/subdir2"),
+                )
+                .unwrap(),
+                children: vec![],
+                ignored: false,
+                loaded: true,
+            };
+
+            // Add the nested file to its subdirectory
+            subdir2_entry
+                .find_or_insert_child(&dirs.tests().join("subdir1/subdir2/nested.txt"))
+                .expect("Should be able to insert into nested subdirectory entry");
+
+            // Build the directory hierarchy
+            subdir1_entry.children.push(Entry::Directory(subdir2_entry));
+
+            expected_directory_entry
+                .children
+                .push(Entry::Directory(subdir1_entry));
+
+            // Create a MerkleNode with the expected final structure
+            let (expected_node, _) = MerkleNode::new(DirEntryOrFragment::Entry(Entry::Directory(
+                expected_directory_entry,
+            )))
+            .expect("Should be able to construct node with expected structure");
+
+            // The hash of our modified node should match the hash of the manually constructed node
+            assert_eq!(
+            node.hash, expected_node.hash,
+            "Hash after all operations should match hash of node constructed with expected structure"
+        );
+        },
+    );
+}

crates/ai/src/index/full_source_code_embedding/merkle_tree/serialized_tree.rs

@@ -0,0 +1,200 @@
+use std::{
+    ops::Range,
+    path::{Path, PathBuf},
+};
+
+use super::{hash::MerkleHash, node::NodeId, MerkleTree, NodeHash, NodeLens};
+
+use crate::index::full_source_code_embedding::fragment_metadata::{
+    FragmentLocation, LeafToFragmentMetadata,
+};
+
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use string_offset::ByteOffset;
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub(crate) struct SerializedCodebaseIndex {
+    tree: SerializedMerkleTree,
+}
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub(crate) struct SerializedMerkleTree {
+    root: SerializedMerkleNode,
+}
+
+/// A given node in the [`MerkleTree`].
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub(super) struct SerializedMerkleNode {
+    /// The hash for the current node of the Merkle tree.
+    pub hash: MerkleHash,
+    /// The children of this merkle node.
+    pub children: Vec<SerializedMerkleNode>,
+    /// Node-specific details.
+    pub fs_info: SerializedFilesystemInfo,
+}
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub(super) struct SerializedFragmentLocation {
+    /// Start line number (inclusive).
+    pub start_line: usize,
+    /// End line number (inclusive).
+    pub end_line: usize,
+    /// The range of byte indices into the original source string for this fragment.
+    pub byte_range: Range<ByteOffset>,
+}
+
+impl From<&FragmentLocation> for SerializedFragmentLocation {
+    fn from(location: &FragmentLocation) -> Self {
+        Self {
+            start_line: location.start_line,
+            end_line: location.end_line,
+            byte_range: location.byte_range.clone(),
+        }
+    }
+}
+
+impl From<&SerializedFragmentLocation> for FragmentLocation {
+    fn from(serialized: &SerializedFragmentLocation) -> Self {
+        Self {
+            start_line: serialized.start_line,
+            end_line: serialized.end_line,
+            byte_range: serialized.byte_range.clone(),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub(super) enum SerializedFilesystemInfo {
+    Directory {
+        absolute_path: PathBuf,
+    },
+    File {
+        /// The path of this node in the filesystem.
+        absolute_path: PathBuf,
+        /// File size in bytes.
+        file_size: usize,
+        /// Time the file was last modified, according to the filesystem.
+        fs_modified_time: DateTime<Utc>,
+        /// SHA-256 hash of the file contents.
+        file_contents_hash: String,
+    },
+    Fragment {
+        /// Fragment location within the file.
+        location: SerializedFragmentLocation,
+    },
+}
+
+impl SerializedCodebaseIndex {
+    pub fn new(
+        tree: &MerkleTree,
+        leaf_node_to_fragment_metadata: &LeafToFragmentMetadata,
+    ) -> anyhow::Result<Self> {
+        let root = SerializedMerkleNode::new(tree.root_node(), leaf_node_to_fragment_metadata)?;
+
+        Ok(Self {
+            tree: SerializedMerkleTree { root },
+        })
+    }
+
+    /// Consumes the index and returns just the tree.
+    pub(crate) fn into_tree(self) -> SerializedMerkleTree {
+        self.tree
+    }
+}
+
+impl SerializedMerkleTree {
+    pub(super) fn into_root(self) -> SerializedMerkleNode {
+        self.root
+    }
+}
+
+fn node_to_filesystem_info(
+    node: &NodeLens,
+    fragment_metadata_mapping: &LeafToFragmentMetadata,
+) -> anyhow::Result<SerializedFilesystemInfo> {
+    let absolute_path = node.path().to_path_buf();
+    match node.node_id() {
+        NodeId::Directory { .. } => Ok(SerializedFilesystemInfo::Directory { absolute_path }),
+        NodeId::File {
+            file_size,
+            fs_modified_time,
+            file_contents_hash,
+            ..
+        } => Ok(SerializedFilesystemInfo::File {
+            absolute_path,
+            file_size: *file_size,
+            fs_modified_time: *fs_modified_time,
+            file_contents_hash: file_contents_hash.clone(),
+        }),
+        NodeId::Fragment {
+            absolute_path,
+            content_range,
+        } => {
+            let Some(metadata_mapping) = fragment_metadata_mapping.get(node.hash()) else {
+                return Err(anyhow::anyhow!(
+                    "did not find hash in fragment metadata mapping"
+                ));
+            };
+
+            let Some(fragment) = metadata_mapping.iter().find(|fragment| {
+                fragment.absolute_path == *absolute_path
+                    && fragment.location.byte_range == *content_range
+            }) else {
+                return Err(anyhow::anyhow!(
+                    "did not find fragment metadata with matching path and content range"
+                ));
+            };
+
+            Ok(SerializedFilesystemInfo::Fragment {
+                location: (&fragment.location).into(),
+            })
+        }
+    }
+}
+
+impl SerializedMerkleNode {
+    fn new(
+        node: NodeLens,
+        fragment_metadata_mapping: &LeafToFragmentMetadata,
+    ) -> anyhow::Result<Self> {
+        let fs_info = node_to_filesystem_info(&node, fragment_metadata_mapping)?;
+
+        let children = node
+            .children()
+            .map(|child| Self::new(child, fragment_metadata_mapping))
+            .collect::<anyhow::Result<Vec<_>>>()?;
+        let hash = node.hash().as_ref().clone();
+        Ok(Self {
+            hash,
+            children,
+            fs_info,
+        })
+    }
+
+    /// Returns the node's absolute path if it is a file or directory.
+    /// Returns None if the node is a fragment.
+    pub(super) fn absolute_path(&self) -> Option<&Path> {
+        match &self.fs_info {
+            SerializedFilesystemInfo::Directory { absolute_path }
+            | SerializedFilesystemInfo::File { absolute_path, .. } => Some(absolute_path.as_path()),
+            SerializedFilesystemInfo::Fragment { .. } => None,
+        }
+    }
+
+    pub(super) fn hash(&self) -> NodeHash {
+        NodeHash::new(self.hash.to_owned())
+    }
+
+    pub(super) fn children(&self) -> impl Iterator<Item = &SerializedMerkleNode> {
+        self.children.iter()
+    }
+
+    pub(super) fn fs_info(&self) -> &SerializedFilesystemInfo {
+        &self.fs_info
+    }
+}
+
+#[cfg(test)]
+#[path = "serialized_tree_test.rs"]
+mod tests;

crates/ai/src/index/full_source_code_embedding/merkle_tree/serialized_tree_test.rs

@@ -0,0 +1,88 @@
+use futures::executor::block_on;
+use serde_json;
+use virtual_fs::VirtualFS;
+
+use crate::index::full_source_code_embedding::merkle_tree::{
+    construct_test_merkle_tree, MerkleTree,
+};
+
+use super::SerializedCodebaseIndex;
+
+#[test]
+fn round_trip_index_serialize_deserialize_json() {
+    VirtualFS::test("test_nodes_from_path_json", |dirs, mut sandbox| {
+        let (original_tree, original_metadata) =
+            block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+        let serializable_index = SerializedCodebaseIndex::new(&original_tree, &original_metadata);
+        let serializable_index =
+            serializable_index.expect("Should successfully construct serializable index");
+
+        let serialized_str =
+            serde_json::to_string(&serializable_index).expect("Should serialize to JSON string");
+        assert!(!serialized_str.is_empty());
+
+        let deserialized_index: SerializedCodebaseIndex =
+            serde_json::from_str(&serialized_str).expect("Should deserialize from JSON");
+        assert_eq!(
+            deserialized_index, serializable_index,
+            "Serialized struct should be identical"
+        );
+
+        let (reconstructed_tree, reconstructed_metadata) =
+            MerkleTree::from_serialized_tree(deserialized_index.into_tree())
+                .expect("Should rebuild Merkle Tree");
+        assert_eq!(
+            reconstructed_tree.root_node().hash(),
+            original_tree.root_node().hash(),
+            "Reconstructed Merkle tree should be identical",
+        );
+        assert_eq!(
+            original_metadata, reconstructed_metadata,
+            "Reconstructed metadata should be identical"
+        );
+    })
+}
+
+#[test]
+fn round_trip_index_serialize_deserialize_bincode() {
+    VirtualFS::test("test_nodes_from_path_bincode", |dirs, mut sandbox| {
+        let (original_tree, original_metadata) =
+            block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+        let serializable_index = SerializedCodebaseIndex::new(&original_tree, &original_metadata);
+        let serializable_index =
+            serializable_index.expect("Should successfully construct serializable index");
+
+        let serialized_bytes =
+            bincode::serialize(&serializable_index).expect("Should serialize to bincode");
+        assert!(!serialized_bytes.is_empty());
+
+        // Bincode output should be smaller than JSON
+        let json_bytes = serde_json::to_vec(&serializable_index).unwrap();
+        assert!(
+            serialized_bytes.len() < json_bytes.len(),
+            "Bincode ({} bytes) should be smaller than JSON ({} bytes)",
+            serialized_bytes.len(),
+            json_bytes.len(),
+        );
+
+        let deserialized_index: SerializedCodebaseIndex =
+            bincode::deserialize(&serialized_bytes).expect("Should deserialize from bincode");
+        assert_eq!(
+            deserialized_index, serializable_index,
+            "Serialized struct should be identical"
+        );
+
+        let (reconstructed_tree, reconstructed_metadata) =
+            MerkleTree::from_serialized_tree(deserialized_index.into_tree())
+                .expect("Should rebuild Merkle Tree");
+        assert_eq!(
+            reconstructed_tree.root_node().hash(),
+            original_tree.root_node().hash(),
+            "Reconstructed Merkle tree should be identical",
+        );
+        assert_eq!(
+            original_metadata, reconstructed_metadata,
+            "Reconstructed metadata should be identical"
+        );
+    })
+}

crates/ai/src/index/full_source_code_embedding/merkle_tree/test_util.rs

@@ -0,0 +1,99 @@
+use super::MerkleTree;
+use crate::index::full_source_code_embedding::fragment_metadata::LeafToFragmentMetadata;
+use repo_metadata::{DirectoryEntry, Entry};
+use virtual_fs::{Stub, VirtualFS};
+
+/// Construct a test Merkle tree with the following structure:
+/// ```
+/// root.txt
+/// top_dir/
+/// ├── file1.txt
+/// ├── subdir_a/
+/// │   ├── file2.txt
+/// │   └── file3.txt
+/// └── subdir_b/
+///     └── file4.txt
+/// ```
+#[cfg(test)]
+pub async fn construct_test_merkle_tree(
+    dirs: &virtual_fs::Dirs,
+    sandbox: &mut VirtualFS,
+) -> (MerkleTree, LeafToFragmentMetadata) {
+    sandbox.mkdir("top_dir");
+    sandbox.mkdir("top_dir/subdir_a");
+    sandbox.mkdir("top_dir/subdir_b");
+    sandbox.with_files(vec![
+        Stub::FileWithContent("root.txt", "root content"),
+        Stub::FileWithContent("top_dir/file1.txt", "file1 content"),
+        Stub::FileWithContent("top_dir/subdir_a/file2.txt", "file2 content"),
+        Stub::FileWithContent("top_dir/subdir_a/file3.txt", "file3 content"),
+        Stub::FileWithContent("top_dir/subdir_b/file4.txt", "file4 content"),
+    ]);
+
+    let mut root_dir_entry = DirectoryEntry {
+        path: warp_util::standardized_path::StandardizedPath::try_from_local(dirs.tests()).unwrap(),
+        children: vec![],
+        ignored: false,
+        loaded: true,
+    };
+    root_dir_entry
+        .find_or_insert_child(&dirs.tests().join("root.txt"))
+        .expect("Should be able to insert root file");
+
+    let mut top_dir_entry = DirectoryEntry {
+        path: warp_util::standardized_path::StandardizedPath::try_from_local(
+            &dirs.tests().join("top_dir"),
+        )
+        .unwrap(),
+        children: vec![],
+        ignored: false,
+        loaded: true,
+    };
+    top_dir_entry
+        .find_or_insert_child(&dirs.tests().join("top_dir/file1.txt"))
+        .expect("Should be able to insert file1");
+
+    let mut subdir_a_entry = DirectoryEntry {
+        path: warp_util::standardized_path::StandardizedPath::try_from_local(
+            &dirs.tests().join("top_dir/subdir_a"),
+        )
+        .unwrap(),
+        children: vec![],
+        ignored: false,
+        loaded: true,
+    };
+    subdir_a_entry
+        .find_or_insert_child(&dirs.tests().join("top_dir/subdir_a/file2.txt"))
+        .expect("Should be able to insert file2");
+    subdir_a_entry
+        .find_or_insert_child(&dirs.tests().join("top_dir/subdir_a/file3.txt"))
+        .expect("Should be able to insert file3");
+
+    let mut subdir_b_entry = DirectoryEntry {
+        path: warp_util::standardized_path::StandardizedPath::try_from_local(
+            &dirs.tests().join("top_dir/subdir_b"),
+        )
+        .unwrap(),
+        children: vec![],
+        ignored: false,
+        loaded: true,
+    };
+    subdir_b_entry
+        .find_or_insert_child(&dirs.tests().join("top_dir/subdir_b/file4.txt"))
+        .expect("Should be able to insert file4");
+
+    top_dir_entry
+        .children
+        .push(Entry::Directory(subdir_a_entry));
+    top_dir_entry
+        .children
+        .push(Entry::Directory(subdir_b_entry));
+
+    root_dir_entry
+        .children
+        .push(Entry::Directory(top_dir_entry));
+
+    MerkleTree::try_new(Entry::Directory(root_dir_entry))
+        .await
+        .expect("Should be able to construct tree")
+}

crates/ai/src/index/full_source_code_embedding/merkle_tree/tree.rs

@@ -0,0 +1,214 @@
+use crate::index::Entry;
+use anyhow::anyhow;
+use cfg_if::cfg_if;
+use std::{
+    collections::{HashSet, VecDeque},
+    path::PathBuf,
+};
+
+use crate::index::full_source_code_embedding::fragment_metadata::{
+    LeafToFragmentMetadata, LeafToFragmentMetadataUpdates,
+};
+use crate::index::full_source_code_embedding::Error;
+
+use super::{
+    node::{ChildrenPath, MerkleNode, NodeLens, NodeMask},
+    serialized_tree::SerializedMerkleTree,
+    DirEntryOrFragment,
+};
+
+pub(super) enum UpdateFileResult {
+    Deleted,
+    Updated,
+    NoChange,
+}
+
+pub(crate) struct TreeUpdateResult<'a> {
+    pub node_lens: Vec<NodeLens<'a>>,
+    pub leaf_to_fragment_meta_updates: LeafToFragmentMetadataUpdates,
+}
+
+/// A merkle tree used for codebase indexing. This data structure allows us to efficiently compute
+/// which parts of a repository have changed without needing to traverse every file in the tree.
+///
+/// The leaves of this tree are code fragments of a given file in the repository, with a
+/// corresponding SHA-256 hash of the contents.
+///
+/// The parents nodes in this tree are either a directory or a file (where the children are all the
+/// fragments of the file) and a corresponding hash of all the hashes of the children.
+///
+/// For example. Consider the following repository structure:
+/// * `/src`
+/// * `/src/foo.rs`
+/// * `/src/bar.rs`
+/// * `/src/bazz/buzz.rs`
+///
+/// The tree would roughly look:
+///
+/// /src (Hash: FooBarBuzzBazz)
+/// ├── /src/foo.rs (Hash: Foo)
+/// ├── /src/bar.rs (Hash: Bar)
+/// └── /src/bazz (Hash: BuzzBazz)
+///     └── /src/bazz/buzz.rs (Hash: Buzz)
+/// `
+#[derive(Debug)]
+pub(crate) struct MerkleTree {
+    root: MerkleNode,
+}
+
+impl MerkleTree {
+    /// Creates a new [`MerkleTree`] given the root node of an [`Entry`].
+    /// Returns an error if a node could not be created for any reason.
+    pub async fn try_new(entry: Entry) -> anyhow::Result<(MerkleTree, LeafToFragmentMetadata)> {
+        let build_node = move || MerkleNode::new(DirEntryOrFragment::Entry(entry));
+
+        let (root, mapping_update) = if tokio::runtime::Handle::try_current().is_ok() {
+            // Offload to a blocking thread so that the rayon `pool.install()` call inside
+            // `MerkleNode::new` does not block a tokio executor thread. Blocking executor
+            // threads starves other async tasks (e.g. shell history parsing during bootstrap).
+            tokio::task::spawn_blocking(build_node)
+                .await
+                .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {e}"))?
+        } else {
+            build_node()
+        }?;
+
+        let leaf_node_to_fragment_metadata = LeafToFragmentMetadata::new(mapping_update);
+        Ok((Self { root }, leaf_node_to_fragment_metadata))
+    }
+
+    pub fn root_node(&self) -> NodeLens<'_> {
+        NodeLens::new(&self.root)
+    }
+
+    pub fn from_serialized_tree(
+        serialized_tree: SerializedMerkleTree,
+    ) -> anyhow::Result<(Self, LeafToFragmentMetadata)> {
+        let serialized_root = serialized_tree.into_root();
+        let Some(root_path) = serialized_root.absolute_path() else {
+            return Err(anyhow::anyhow!("root node should never be a fragment"));
+        };
+        let root_path = root_path.to_path_buf();
+        let (root, mapping_update) =
+            MerkleNode::from_serialized(serialized_root, root_path.as_path())?;
+        let leaf_node_to_fragment_metadata = LeafToFragmentMetadata::new(mapping_update);
+        Ok((Self { root }, leaf_node_to_fragment_metadata))
+    }
+
+    /// Construct the changed nodes' NodeLens from a NodeMask.
+    /// NodeLens are returned in reverse-BFS order (children first).
+    pub(super) fn nodes_from_mask(&self, node_mask: NodeMask) -> Result<Vec<NodeLens<'_>>, Error> {
+        let mut result = vec![];
+        let mut queue = VecDeque::new();
+        queue.push_back((&self.root, node_mask));
+
+        while let Some((current_node, path)) = queue.pop_front() {
+            result.push(NodeLens::new(current_node));
+
+            match path.children {
+                ChildrenPath::All => {
+                    for (idx, child_node) in current_node.children().enumerate() {
+                        queue.push_back((child_node, NodeMask::new(idx)));
+                    }
+                }
+                ChildrenPath::SpecificChildren(children_path) => {
+                    for child_path in children_path {
+                        if child_path.index < current_node.count_children() {
+                            let child_node = &current_node.child_at(child_path.index);
+                            queue.push_back((child_node, child_path));
+                        } else {
+                            return Err(Error::Other(anyhow!(
+                                "Invalid child index {} in node with {} child(ren)",
+                                child_path.index,
+                                current_node.count_children(),
+                            )));
+                        }
+                    }
+                }
+            }
+        }
+
+        result.reverse();
+        Ok(result)
+    }
+
+    /// Given the paths to a set of changed files, for each file, update the node if it exists in the tree, or
+    /// create a node if it doesn't exist. This also updates all the intermediate nodes.
+    ///
+    /// Return all the updated nodes.
+    pub async fn upsert_files(
+        &mut self,
+        mut paths: HashSet<PathBuf>,
+    ) -> Result<TreeUpdateResult<'_>, Error> {
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_meta_updates = LeafToFragmentMetadataUpdates::empty();
+
+        let mut do_upsert = || {
+            self.root.upsert_files(
+                &mut paths,
+                &mut node_path,
+                &mut leaf_to_fragment_meta_updates,
+            )
+        };
+
+        cfg_if! {
+            if #[cfg(not(target_family = "wasm"))] {
+                let upsert_result = if tokio::runtime::Handle::try_current().is_ok() {
+                    // `upsert_files` is expensive and can block a background thread for a while,
+                    // so use `block_in_place` to tell tokio to move any tasks enqueued for this
+                    // thread to another thread (so that they might be able to run on a different
+                    // thread).
+                    tokio::task::block_in_place(do_upsert)
+                } else {
+                    do_upsert()
+                };
+            } else {
+                let upsert_result = do_upsert();
+            }
+        }
+
+        // Note that we cannot early return directly here on error. We need to make sure leaf_node_to_fragment_metadatas
+        // is properly written so the tree remains valid.
+        let node_lens = if !matches!(upsert_result, UpdateFileResult::NoChange) {
+            self.nodes_from_mask(node_path)?
+        } else {
+            vec![self.root_node()]
+        };
+
+        Ok(TreeUpdateResult {
+            node_lens,
+            leaf_to_fragment_meta_updates,
+        })
+    }
+
+    /// Given the path to a removed file, remove the node if it exists in the tree.
+    ///
+    /// Return all the updated nodes.
+    pub async fn remove_files(
+        &mut self,
+        mut paths: HashSet<PathBuf>,
+    ) -> Result<TreeUpdateResult<'_>, Error> {
+        let mut node_path = NodeMask::default();
+        let mut leaf_to_fragment_meta_updates = LeafToFragmentMetadataUpdates::empty();
+
+        // Note that we cannot early return directly here on error. We need to make sure leaf_node_to_fragment_metadatas
+        // is properly written so the tree remains valid.
+        let node_lens = match self.root.remove_files(
+            &mut paths,
+            &mut node_path,
+            &mut leaf_to_fragment_meta_updates,
+        ) {
+            UpdateFileResult::NoChange => vec![self.root_node()],
+            _ => self.nodes_from_mask(node_path)?,
+        };
+
+        Ok(TreeUpdateResult {
+            node_lens,
+            leaf_to_fragment_meta_updates,
+        })
+    }
+}
+
+#[cfg(test)]
+#[path = "tree_test.rs"]
+mod tests;

crates/ai/src/index/full_source_code_embedding/merkle_tree/tree_test.rs

@@ -0,0 +1,121 @@
+use crate::index::full_source_code_embedding::merkle_tree::{
+    construct_test_merkle_tree, node::ChildrenPath,
+};
+use futures::executor::block_on;
+
+use virtual_fs::VirtualFS;
+
+use super::*;
+
+#[test]
+fn test_nodes_from_path() {
+    VirtualFS::test("test_nodes_from_path", |dirs, mut sandbox| {
+        let (tree, _metadata) = block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+
+        // Test: path with single child
+        let single_path = NodeMask {
+            index: 0,
+            children: ChildrenPath::SpecificChildren(vec![NodeMask::new(0)]),
+        };
+        let result = tree.nodes_from_mask(single_path).expect("Should not fail");
+        assert_eq!(
+            result.len(),
+            3,
+            "Should return three nodes for single-level path: root, child and its children"
+        );
+
+        // Test: path with multiple levels
+        // top_dir/file1.txt
+        let multi_path = NodeMask {
+            index: 0,
+            children: ChildrenPath::SpecificChildren(vec![NodeMask {
+                index: 1,
+                children: ChildrenPath::SpecificChildren(vec![NodeMask {
+                    index: 1,
+                    children: ChildrenPath::SpecificChildren(vec![NodeMask::new(0)]),
+                }]),
+            }]),
+        };
+        let result = tree.nodes_from_mask(multi_path).expect("Should not fail");
+        assert_eq!(
+            result.len(),
+            4,
+            "Should return 4 nodes for multi-level path: root, top_dir, file1.txt, and file1.txt's contents"
+        );
+        assert!(
+            result.first().unwrap().is_leaf(),
+            "First node should be file1.txt's contents"
+        );
+
+        // Test: invalid index should return error
+        let invalid_path = NodeMask {
+            index: 0,
+            children: ChildrenPath::SpecificChildren(vec![NodeMask::new(999)]),
+        };
+        assert!(
+            tree.nodes_from_mask(invalid_path).is_err(),
+            "Should return error for invalid index"
+        );
+
+        // Test: verify nodes are returned in reverse BFS order (children first)
+        let path = NodeMask {
+            index: 0,
+            children: ChildrenPath::SpecificChildren(vec![
+                NodeMask {
+                    // root.txt
+                    index: 0,
+                    children: ChildrenPath::SpecificChildren(vec![NodeMask::new(0)]),
+                },
+                NodeMask {
+                    // top_dir
+                    index: 1,
+                    children: ChildrenPath::SpecificChildren(vec![
+                        NodeMask {
+                            // subdir_b
+                            index: 0,
+                            children: ChildrenPath::SpecificChildren(vec![NodeMask {
+                                // file4.txt
+                                index: 0,
+                                children: ChildrenPath::SpecificChildren(vec![NodeMask::new(0)]),
+                            }]),
+                        },
+                        NodeMask::new(1), // file1.txt
+                        NodeMask {
+                            // subdir_a
+                            index: 2,
+                            children: ChildrenPath::SpecificChildren(vec![
+                                NodeMask {
+                                    // file2.txt
+                                    index: 0,
+                                    children: ChildrenPath::SpecificChildren(vec![NodeMask::new(
+                                        0,
+                                    )]),
+                                },
+                                NodeMask {
+                                    // file3.txt
+                                    index: 1,
+                                    children: ChildrenPath::SpecificChildren(vec![NodeMask::new(
+                                        0,
+                                    )]),
+                                },
+                            ]),
+                        },
+                    ]),
+                },
+            ]),
+        };
+
+        let result = tree.nodes_from_mask(path).expect("Should not fail");
+
+        // Children should come first
+        assert!(result[0].is_leaf(), "First node should be a leaf");
+        assert!(result[1].is_leaf(), "Second node should be a leaf");
+        assert!(result[2].is_leaf(), "Third node should be a leaf");
+
+        // Last node should be root
+        assert!(
+            !result.last().unwrap().is_leaf(),
+            "Last node should not be a leaf"
+        );
+    });
+}

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

@@ -0,0 +1,203 @@
+mod changed_files;
+mod chunker;
+mod codebase_index;
+mod fragment_metadata;
+pub mod manager;
+mod merkle_tree;
+mod priority_queue;
+mod snapshot;
+pub mod store_client;
+mod sync_client;
+
+use std::{ops::Range, path::PathBuf, time::Duration};
+pub use sync_client::SyncTask;
+
+pub use codebase_index::{CodebaseIndex, RetrievalID, SyncProgress};
+pub use merkle_tree::{ContentHash, NodeHash};
+
+use fragment_metadata::FragmentMetadata;
+use string_offset::ByteOffset;
+use thiserror::Error;
+use warp_graphql::queries::rerank_fragments::FragmentLocationInput;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("File I/O error {0:#}")]
+    Io(#[from] std::io::Error),
+    #[error("Not a git repository")]
+    NotAGitRepository,
+    #[error("Build tree error {0:#}")]
+    BuildTreeError(#[from] crate::index::BuildTreeError),
+    #[error("Unsupported platform")]
+    UnsupportedPlatform,
+    #[error("Invalid hash: {0:#}")]
+    InvalidHash(base16ct::Error),
+    #[error("Empty node content")]
+    EmptyNodeContent,
+    #[error("Failed to get metadata")]
+    FailedToGetMetadata(PathBuf),
+    #[error("File size exceeds maximum limit")]
+    FileSizeExceeded,
+    #[error(transparent)]
+    InconsistentState(#[from] InconsistentStateError),
+    #[error("Failed to generate embeddings for some hashes")]
+    FailedToGenerateEmbeddings(Vec<FragmentMetadata>),
+    #[error("Failed to sync some intermediate nodes")]
+    FailedToSyncIntermediateNodes(Vec<NodeHash>),
+    #[error("Diff merkle tree {0:#}")]
+    DiffMerkleTreeError(#[from] crate::index::full_source_code_embedding::DiffMerkleTreeError),
+    #[error("File system changed since merkle tree construction")]
+    FileSystemStateChanged,
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+    #[error("Failed to parse snapshot")]
+    SnapshotParsingFailed,
+}
+
+// Based off of BuildTreeError in entry.rs
+#[derive(Debug, Error)]
+pub enum DiffMerkleTreeError {
+    #[error("Merkle tree node and file mismatch")]
+    CurrentNodeMismatch(PathBuf),
+    #[error("File is ignored")]
+    Ignored,
+    #[error("Symlink is not supported")]
+    Symlink,
+    #[error("Fragment node in diffing process")]
+    Fragment(PathBuf),
+    #[error("Max depth exceeded")]
+    MaxDepthExceeded,
+    #[error("Exceeded max file limit")]
+    ExceededMaxFileLimit,
+}
+
+#[derive(Error, Debug)]
+pub enum InconsistentStateError {
+    #[error("Missing fragment metadata for {fragment_hash}")]
+    MissingFragmentMetadata { fragment_hash: ContentHash },
+    #[error("Can't find node index in merkle node")]
+    NodeIndexNotFound,
+}
+
+#[allow(non_camel_case_types)]
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub enum EmbeddingConfig {
+    OpenAiTextSmall3_256,
+    VoyageCode3_512,
+    Voyage3_5_Lite_512,
+    #[default]
+    Voyage3_5_512,
+}
+
+#[derive(Debug, Clone)]
+pub struct RepoMetadata {
+    pub path: Option<String>,
+}
+
+impl From<RepoMetadata> for warp_graphql::full_source_code_embedding::RepoMetadata {
+    fn from(val: RepoMetadata) -> Self {
+        Self { path: val.path }
+    }
+}
+
+impl From<EmbeddingConfig> for warp_graphql::full_source_code_embedding::EmbeddingConfig {
+    fn from(val: EmbeddingConfig) -> Self {
+        match val {
+            EmbeddingConfig::OpenAiTextSmall3_256 => {
+                warp_graphql::full_source_code_embedding::EmbeddingConfig::OpenaiTextSmall3256
+            }
+            EmbeddingConfig::VoyageCode3_512 => {
+                warp_graphql::full_source_code_embedding::EmbeddingConfig::VoyageCode3512
+            }
+            EmbeddingConfig::Voyage3_5_512 => {
+                warp_graphql::full_source_code_embedding::EmbeddingConfig::Voyage35512
+            }
+            EmbeddingConfig::Voyage3_5_Lite_512 => {
+                warp_graphql::full_source_code_embedding::EmbeddingConfig::Voyage35Lite512
+            }
+        }
+    }
+}
+
+impl TryFrom<warp_graphql::full_source_code_embedding::EmbeddingConfig> for EmbeddingConfig {
+    type Error = Error;
+
+    fn try_from(
+        value: warp_graphql::full_source_code_embedding::EmbeddingConfig,
+    ) -> Result<Self, Self::Error> {
+        match value {
+            warp_graphql::full_source_code_embedding::EmbeddingConfig::OpenaiTextSmall3256 => {
+                Ok(Self::OpenAiTextSmall3_256)
+            }
+            warp_graphql::full_source_code_embedding::EmbeddingConfig::Voyage35Lite512 => {
+                Ok(Self::Voyage3_5_Lite_512)
+            }
+            warp_graphql::full_source_code_embedding::EmbeddingConfig::VoyageCode3512 => {
+                Ok(Self::VoyageCode3_512)
+            }
+            warp_graphql::full_source_code_embedding::EmbeddingConfig::Voyage35512 => {
+                Ok(Self::Voyage3_5_512)
+            }
+        }
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct CodebaseContextConfig {
+    pub embedding_config: EmbeddingConfig,
+    pub embedding_cadence: Duration,
+}
+
+#[derive(Clone)]
+pub struct FragmentLocation {
+    absolute_path: PathBuf,
+    byte_range: Range<ByteOffset>,
+}
+
+#[derive(Clone)]
+pub struct Fragment {
+    content: String,
+    content_hash: ContentHash,
+    location: FragmentLocation,
+}
+
+impl From<Fragment> for warp_graphql::full_source_code_embedding::Fragment {
+    fn from(val: Fragment) -> Self {
+        Self {
+            content: val.content,
+            content_hash: val.content_hash.into(),
+        }
+    }
+}
+
+impl From<Fragment> for warp_graphql::queries::rerank_fragments::RerankFragmentInput {
+    fn from(val: Fragment) -> Self {
+        Self {
+            content: val.content,
+            content_hash: val.content_hash.into(),
+            location: FragmentLocationInput {
+                byte_start: val.location.byte_range.start.as_usize() as i32,
+                byte_end: val.location.byte_range.end.as_usize() as i32,
+                file_path: val.location.absolute_path.to_string_lossy().to_string(),
+            },
+        }
+    }
+}
+
+impl TryFrom<warp_graphql::queries::rerank_fragments::RerankFragment> for Fragment {
+    type Error = Error;
+
+    fn try_from(
+        val: warp_graphql::queries::rerank_fragments::RerankFragment,
+    ) -> Result<Self, Self::Error> {
+        Ok(Self {
+            content: val.content,
+            content_hash: val.content_hash.try_into()?,
+            location: FragmentLocation {
+                absolute_path: PathBuf::from(val.location.file_path),
+                byte_range: ByteOffset::from(val.location.byte_start as usize)
+                    ..ByteOffset::from(val.location.byte_end as usize),
+            },
+        })
+    }
+}

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

@@ -0,0 +1,79 @@
+use std::hash::Hash;
+use std::path::PathBuf;
+
+use itertools::Itertools;
+use priority_queue::PriorityQueue;
+
+use crate::workspace::WorkspaceMetadata;
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub(super) enum Priority {
+    ActiveSession = 0,
+    OpenSession = 1,
+    PersistedSnapshot = 2,
+}
+
+#[derive(Debug, Clone)]
+struct QueueEntry {
+    metadata: WorkspaceMetadata,
+}
+
+impl Hash for QueueEntry {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.metadata.path.hash(state);
+    }
+}
+
+impl PartialEq for QueueEntry {
+    fn eq(&self, other: &Self) -> bool {
+        self.metadata.path == other.metadata.path
+    }
+}
+
+impl Eq for QueueEntry {}
+
+#[derive(Debug, Default)]
+pub(super) struct BuildQueue {
+    queue: PriorityQueue<QueueEntry, Priority>,
+}
+
+impl BuildQueue {
+    pub(super) fn empty() -> Self {
+        Self::default()
+    }
+
+    pub(super) fn queued_metadata(&self) -> impl IntoIterator<Item = WorkspaceMetadata> + use<'_> {
+        self.queue.iter().map(|(entry, _)| entry.metadata.clone())
+    }
+
+    pub(super) fn new_with_persisted(snapshots_to_load: Vec<WorkspaceMetadata>) -> 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)),
+        );
+
+        Self { queue }
+    }
+
+    /// Pulls the next index root path to sync from the priority queue and returns it.
+    pub fn pick_next_sync(&mut self) -> Option<WorkspaceMetadata> {
+        self.queue.pop().map(|(entry, _priority)| entry.metadata)
+    }
+
+    /// Adjusts the priority of a path in the queue if it exists.
+    pub(super) fn update_path_priority(&mut self, root_path: PathBuf, priority: Priority) {
+        // Exemplar is only used to lookup the item in the queue with the Eq implemented above
+        // It will not overwrite the found item.
+        let exemplar = QueueEntry {
+            metadata: WorkspaceMetadata {
+                path: root_path,
+                ..Default::default()
+            },
+        };
+
+        self.queue.change_priority(&exemplar, priority);
+    }
+}

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

@@ -0,0 +1,230 @@
+use chrono::Utc;
+#[cfg(feature = "local_fs")]
+use repo_metadata::Repository;
+use std::{
+    collections::HashSet,
+    hash::{DefaultHasher, Hash, Hasher},
+    path::{Path, PathBuf},
+    time::Duration,
+};
+#[cfg(feature = "local_fs")]
+use warpui::ModelHandle;
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "local_fs")] {
+        use super::Error as CodebaseIndexError;
+        use std::sync::Arc;
+        use warpui::ModelContext;
+        use anyhow::Context;
+        use warp_core::safe_info;
+        use super::{store_client::StoreClient, CodebaseIndex, EmbeddingConfig};
+    }
+}
+
+use crate::workspace::WorkspaceMetadata;
+
+/// Number of days after which an index snapshot should be considered expired.
+pub(super) const REPO_SNAPSHOT_SHELF_LIFE_DAYS: u64 = 30;
+
+/// The maximum lifetime of an index snapshot file, after which it
+/// should be considered expired and deleted.
+const REPO_SNAPSHOT_SHELF_LIFE_DURATION: Duration =
+    Duration::from_secs(60 * 60 * 24 * REPO_SNAPSHOT_SHELF_LIFE_DAYS);
+
+/// Subdirectory inside the app's statedirectory that holds snapshot files.
+const REPO_SNAPSHOT_SUBDIR_NAME: &str = "codebase_index_snapshots";
+
+/// Splits a list of codebase indices into invalid and valid indices,
+/// based on their last write date and whether they have a corresponding snapshot file.
+pub(super) fn split_snapshot_metadata_by_validity(
+    persisted_codebase_indices: Vec<WorkspaceMetadata>,
+) -> (Vec<WorkspaceMetadata>, Vec<WorkspaceMetadata>) {
+    let now = Utc::now();
+    persisted_codebase_indices
+        .into_iter()
+        .partition(|index_metadata| {
+            log::info!(
+                "Discarding expired codebase index snapshot for {:?}",
+                index_metadata.path
+            );
+            index_metadata.is_expired(now, REPO_SNAPSHOT_SHELF_LIFE_DAYS)
+                || !has_snapshot(&index_metadata.path)
+        })
+}
+
+/// Delete snapshot files that are missing metadata or have expired from the snapshot directory.
+pub(super) fn clean_up_snapshot_files(
+    snapshot_file_dir: &Path,
+    persisted_codebase_indices: &[WorkspaceMetadata],
+) {
+    let expected_snapshot_filenames: HashSet<_> = persisted_codebase_indices
+        .iter()
+        .map(|index_metadata| snapshot_path(snapshot_file_dir, &index_metadata.path))
+        .collect();
+
+    if let Ok(fs_entries) = std::fs::read_dir(snapshot_file_dir) {
+        let fs_now = std::time::SystemTime::now();
+        for fs_entry in fs_entries.flatten() {
+            let path = fs_entry.path();
+
+            // Check if this is a regular file with a snapshot_ prefix
+            if let Ok(fs_metadata) = fs_entry.metadata() {
+                if fs_metadata.is_file() {
+                    maybe_clean_up_snapshot_file(
+                        &path,
+                        &fs_metadata,
+                        &fs_now,
+                        &expected_snapshot_filenames,
+                    );
+                }
+            }
+        }
+    }
+}
+
+fn maybe_clean_up_snapshot_file(
+    path: &Path,
+    fs_metadata: &std::fs::Metadata,
+    fs_now: &std::time::SystemTime,
+    expected_snapshot_filenames: &HashSet<PathBuf>,
+) {
+    if let Some(filename) = path.file_name() {
+        if filename.to_string_lossy().starts_with("snapshot_") {
+            let mut should_remove = false;
+
+            // Check if file itself is expired
+            if let Ok(modified_time) = fs_metadata.modified() {
+                if let Ok(age) = fs_now.duration_since(modified_time) {
+                    if age >= REPO_SNAPSHOT_SHELF_LIFE_DURATION {
+                        should_remove = true;
+                    }
+                }
+            }
+
+            // Check if file is not in alive_files set
+            if !expected_snapshot_filenames.contains(path) {
+                should_remove = true;
+            }
+
+            // Remove file if either condition is true
+            if should_remove {
+                if let Err(e) = std::fs::remove_file(path) {
+                    log::warn!("Failed to remove stale snapshot file {path:?}: {e}");
+                }
+            }
+        }
+    }
+}
+
+#[cfg(feature = "local_fs")]
+pub(super) fn read_snapshot(
+    store_client: Arc<dyn StoreClient>,
+    snapshot_dir: &Path,
+    repository: ModelHandle<Repository>,
+    max_files_repo_limit: usize,
+    embedding_generation_batch_size: usize,
+    ctx: &mut ModelContext<CodebaseIndex>,
+) -> anyhow::Result<CodebaseIndex> {
+    let repo_path_buf = repository.as_ref(ctx).root_dir().to_local_path_lossy();
+    let repo_path = repo_path_buf.as_path();
+    let snapshot_path = snapshot_path(snapshot_dir, repo_path);
+    let snapshot_bytes = std::fs::read(&snapshot_path)?;
+
+    let result = CodebaseIndex::new_from_snapshot(
+        repository,
+        store_client.clone(),
+        EmbeddingConfig::default(),
+        snapshot_bytes,
+        max_files_repo_limit,
+        embedding_generation_batch_size,
+        ctx,
+    );
+
+    // If rebuilding merkle tree from snapshot fails due to parsing error, delete the snapshot file.
+    if let Err(CodebaseIndexError::SnapshotParsingFailed) = result {
+        log::info!(
+            "Deleting invalid snapshot {:?} for repo",
+            snapshot_path.display(),
+        );
+        if let Err(e) = std::fs::remove_file(&snapshot_path) {
+            log::warn!(
+                "Failed to remove invalid snapshot file {:?}: {}",
+                snapshot_path.display(),
+                e
+            );
+        }
+    }
+
+    Ok(result?)
+}
+
+pub(super) fn has_snapshot(repo_path: &Path) -> bool {
+    let Some(snapshot_dir) = snapshot_dir() else {
+        return false;
+    };
+
+    let snapshot_path = snapshot_path(snapshot_dir.as_path(), repo_path);
+    snapshot_path.is_file()
+}
+
+/// Construct a directory to store index snapshots, if it doesn't already exist,
+/// and return its path.
+pub(super) fn snapshot_dir() -> Option<PathBuf> {
+    #[cfg(not(feature = "local_fs"))]
+    return None;
+
+    #[cfg(feature = "local_fs")]
+    {
+        let base_dir =
+            warp_core::paths::secure_state_dir().unwrap_or_else(warp_core::paths::state_dir);
+        let snapshot_dir_path = base_dir.join(REPO_SNAPSHOT_SUBDIR_NAME);
+
+        if !snapshot_dir_path.is_dir() {
+            std::fs::create_dir_all(&snapshot_dir_path).ok()?;
+        }
+        Some(snapshot_dir_path)
+    }
+}
+
+/// Constructs a snapshot path given a base directory and the codebase index's root path.
+pub(super) fn snapshot_path(snapshot_dir: &Path, repo_path: &Path) -> PathBuf {
+    // Use a hash the repo_path to create a unique filename
+    let mut hasher = DefaultHasher::new();
+    repo_path.hash(&mut hasher);
+    let snapshot_file_name = format!("snapshot_{}", hasher.finish());
+    snapshot_dir.join(snapshot_file_name)
+}
+
+#[cfg(test)]
+#[path = "snapshot_tests.rs"]
+mod tests;
+
+#[cfg(feature = "local_fs")]
+pub(super) fn migrate_snapshots_to_secure_dir_if_needed() -> anyhow::Result<()> {
+    // Only perform migration if a secure state directory is available.
+    let Some(secure_base) = warp_core::paths::secure_state_dir() else {
+        return Ok(());
+    };
+
+    let new_dir = secure_base.join(REPO_SNAPSHOT_SUBDIR_NAME);
+    let old_dir = warp_core::paths::state_dir().join(REPO_SNAPSHOT_SUBDIR_NAME);
+
+    if new_dir == old_dir {
+        return Ok(());
+    }
+
+    if old_dir.exists() && !new_dir.exists() {
+        if let Some(parent) = new_dir.parent() {
+            std::fs::create_dir_all(parent)
+                .context("Failed to create application data directory")?;
+        }
+        std::fs::rename(&old_dir, &new_dir)
+            .context("Failed to migrate codebase index snapshots")?;
+        safe_info!(
+            safe: ("Migrated codebase index snapshots into secure application container"),
+            full: ("Migrated codebase index snapshots from `{}` to `{}`", old_dir.display(), new_dir.display())
+        );
+    }
+
+    Ok(())
+}

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

@@ -0,0 +1,107 @@
+use chrono::Duration;
+use virtual_fs::{Stub, VirtualFS};
+
+use super::*;
+
+#[test]
+fn test_clean_up_snapshot_files() {
+    VirtualFS::test("test_clean_up_snapshot_files", |dirs, mut sandbox| {
+        // Create snapshot directory with test files
+        sandbox.mkdir(REPO_SNAPSHOT_SUBDIR_NAME);
+
+        // Create a valid snapshot file
+        let test_path = PathBuf::from("/test/path");
+        let mut hasher = DefaultHasher::new();
+        test_path.hash(&mut hasher);
+        let valid_snapshot_name = format!("snapshot_{}", hasher.finish());
+
+        // Create entries in the virtual filesystem
+        let mut snapshot_dir_relative_path = PathBuf::new();
+        snapshot_dir_relative_path.push(REPO_SNAPSHOT_SUBDIR_NAME);
+        sandbox.with_files(vec![
+            // Valid snapshot file that matches metadata
+            Stub::FileWithContent(
+                snapshot_dir_relative_path
+                    .join(&valid_snapshot_name)
+                    .to_string_lossy()
+                    .as_ref(),
+                "valid content",
+            ),
+            // Expired snapshot file
+            Stub::FileWithContent(
+                snapshot_dir_relative_path
+                    .join("snapshot_expired")
+                    .to_string_lossy()
+                    .as_ref(),
+                "expired content",
+            ),
+            // Non-snapshot file that should be ignored
+            Stub::FileWithContent(
+                snapshot_dir_relative_path
+                    .join("regular_file.txt")
+                    .to_string_lossy()
+                    .as_ref(),
+                "regular content",
+            ),
+        ]);
+        // Subdirectory with the 'snapshot_' prefix that should be ignored
+        let invalid_subdir_path = snapshot_dir_relative_path.join("snapshot_prefixed_directory");
+        sandbox.mkdir(invalid_subdir_path.to_string_lossy().to_string().as_str());
+
+        let snapshot_dir_absolute_path = dirs.tests().join(REPO_SNAPSHOT_SUBDIR_NAME);
+
+        let valid_snapshot_file = snapshot_dir_absolute_path.join(&valid_snapshot_name);
+        let expired_file = snapshot_dir_absolute_path.join("snapshot_expired");
+        let regular_file = snapshot_dir_absolute_path.join("regular_file.txt");
+        let prefixed_subdir = snapshot_dir_absolute_path.join("snapshot_prefixed_directory");
+
+        assert!(valid_snapshot_file.is_file());
+        assert!(expired_file.is_file());
+        assert!(regular_file.is_file());
+        assert!(prefixed_subdir.is_dir());
+
+        // Set the expired file to be older than shelf life
+        let old_time = std::time::SystemTime::now()
+            - REPO_SNAPSHOT_SHELF_LIFE_DURATION
+            - Duration::days(1).to_std().unwrap();
+        filetime::set_file_mtime(
+            &expired_file,
+            filetime::FileTime::from_system_time(old_time),
+        )
+        .unwrap();
+
+        // Create test metadata that only includes the valid file
+        let metadata = vec![WorkspaceMetadata {
+            path: test_path,
+            navigated_ts: None,
+            modified_ts: None,
+            queried_ts: None,
+        }];
+
+        // Run cleanup
+        clean_up_snapshot_files(&snapshot_dir_absolute_path, &metadata);
+
+        // Valid snapshot should still exist
+        assert!(valid_snapshot_file.exists());
+
+        // Expired snapshot should be deleted
+        assert!(!expired_file.exists());
+
+        // Regular file should remain untouched
+        assert!(regular_file.exists());
+
+        // Subdirectory with 'snapshot_' prefix should remain untouched
+        assert!(prefixed_subdir.exists());
+    });
+}
+
+#[test]
+fn test_clean_up_snapshot_files_no_snapshot_dir() {
+    VirtualFS::test("test_clean_up_snapshot_files_no_dir", |dirs, _sandbox| {
+        // Test with empty metadata when snapshot directory doesn't exist
+        let snapshot_metadata = vec![];
+        let snapshot_directory = dirs.tests().join(REPO_SNAPSHOT_SUBDIR_NAME);
+        clean_up_snapshot_files(&snapshot_directory, &snapshot_metadata);
+        assert!(snapshot_metadata.is_empty());
+    });
+}

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

@@ -0,0 +1,138 @@
+use async_trait::async_trait;
+use std::{
+    collections::{HashMap, HashSet},
+    fmt::Debug,
+    time::Duration,
+};
+
+use super::{
+    CodebaseContextConfig, ContentHash, EmbeddingConfig, Error, Fragment, NodeHash, RepoMetadata,
+};
+
+/// Client interface for a remote full source code embedding store.
+#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
+pub trait StoreClient: 'static + Send + Sync {
+    /// Persist new intermediate Merkle tree nodes.
+    async fn update_intermediate_nodes(
+        &self,
+        embedding_config: EmbeddingConfig,
+        nodes: Vec<IntermediateNode>,
+    ) -> Result<HashMap<NodeHash, bool>, Error>;
+
+    /// Generate embeddings for individual code fragments.
+    ///
+    /// Embedding generation may fail on a per-fragment basis, so this returns the status of each
+    /// fragment. If the overall request fails, assume that no embeddings were generated.
+    async fn generate_embeddings(
+        &self,
+        embedding_config: EmbeddingConfig,
+        fragments: Vec<Fragment>,
+        root_hash: NodeHash,
+        repo_metadata: RepoMetadata,
+    ) -> Result<HashMap<ContentHash, bool>, Error>;
+
+    async fn populate_merkle_tree_cache(
+        &self,
+        embedding_config: EmbeddingConfig,
+        root_hash: NodeHash,
+        repo_metadata: RepoMetadata,
+    ) -> Result<bool, Error>;
+
+    async fn sync_merkle_tree(
+        &self,
+        nodes: Vec<NodeHash>,
+        embedding_config: EmbeddingConfig,
+    ) -> Result<HashSet<NodeHash>, Error>;
+
+    async fn rerank_fragments(
+        &self,
+        query: String,
+        fragment: Vec<Fragment>,
+    ) -> Result<Vec<Fragment>, Error>;
+
+    async fn get_relevant_fragments(
+        &self,
+        embedding_config: EmbeddingConfig,
+        query: String,
+        root_hash: NodeHash,
+        repo_metadata: RepoMetadata,
+    ) -> Result<Vec<ContentHash>, Error>;
+
+    async fn codebase_context_config(&self) -> Result<CodebaseContextConfig, Error>;
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct MockStoreClient;
+
+#[cfg_attr(not(target_family = "wasm"), async_trait)]
+#[cfg_attr(target_family = "wasm", async_trait(?Send))]
+impl StoreClient for MockStoreClient {
+    async fn update_intermediate_nodes(
+        &self,
+        _embedding_config: EmbeddingConfig,
+        _nodes: Vec<IntermediateNode>,
+    ) -> Result<HashMap<NodeHash, bool>, Error> {
+        Ok(HashMap::new())
+    }
+
+    async fn generate_embeddings(
+        &self,
+        _embedding_config: EmbeddingConfig,
+        _fragments: Vec<Fragment>,
+        _root_hash: NodeHash,
+        _repo_metadata: RepoMetadata,
+    ) -> Result<HashMap<ContentHash, bool>, Error> {
+        Ok(HashMap::new())
+    }
+
+    async fn sync_merkle_tree(
+        &self,
+        _nodes: Vec<NodeHash>,
+        _embedding_config: EmbeddingConfig,
+    ) -> Result<HashSet<NodeHash>, Error> {
+        Ok(HashSet::new())
+    }
+
+    async fn populate_merkle_tree_cache(
+        &self,
+        _embedding_config: EmbeddingConfig,
+        _root_hash: NodeHash,
+        _repo_metadata: RepoMetadata,
+    ) -> Result<bool, Error> {
+        Ok(true)
+    }
+
+    async fn rerank_fragments(
+        &self,
+        _query: String,
+        fragments: Vec<Fragment>,
+    ) -> Result<Vec<Fragment>, Error> {
+        // Return input as is for mock
+        Ok(fragments)
+    }
+
+    async fn get_relevant_fragments(
+        &self,
+        _embedding_config: EmbeddingConfig,
+        _query: String,
+        _root_hash: NodeHash,
+        _repo_metadata: RepoMetadata,
+    ) -> Result<Vec<ContentHash>, Error> {
+        Ok(Vec::new())
+    }
+
+    async fn codebase_context_config(&self) -> Result<CodebaseContextConfig, Error> {
+        Ok(CodebaseContextConfig {
+            embedding_config: EmbeddingConfig::default(),
+            embedding_cadence: Duration::from_secs(300),
+        })
+    }
+}
+
+/// The contents of an intermediate Merkle tree node, used to sync it to the remote embedding store.
+#[derive(Debug, Clone)]
+pub struct IntermediateNode {
+    pub hash: NodeHash,
+    pub children: Vec<NodeHash>,
+}

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

@@ -0,0 +1,604 @@
+use anyhow::{anyhow, Result};
+use itertools::Itertools;
+use std::future::Future;
+use std::ops::AddAssign;
+use std::pin::Pin;
+use std::{
+    collections::{HashMap, HashSet},
+    mem,
+    sync::Arc,
+};
+use warp_core::sync_queue::{IsTransientError, SyncQueue, SyncQueueTaskTrait};
+
+use super::{CodebaseContextConfig, NodeHash};
+
+use crate::index::full_source_code_embedding::store_client::IntermediateNode;
+
+use super::{
+    changed_files::ChangedFiles,
+    codebase_index::{build_fragments_from_metadata, SyncProgress},
+    fragment_metadata::LeafToFragmentMetadataMapping,
+    merkle_tree::{MerkleTree, NodeLens},
+    store_client::StoreClient,
+    EmbeddingConfig, Error, RepoMetadata,
+};
+use super::{ContentHash, Fragment};
+
+const SYNC_NODE_BATCH_SIZE: usize = 500;
+// Minimum node batch size used for updates.
+const MIN_UPDATE_NODE_BATCH_SIZE: usize = 100;
+
+/// Maximum total raw content bytes per `GenerateCodeEmbeddings` request.
+/// Set to 4 MB to stay under the 5 MB Cloud Armor limit after JSON serialization overhead.
+const MAX_BATCH_CONTENT_BYTES: usize = 4_000_000;
+
+#[derive(Debug, Clone, Default)]
+pub struct FlushFragmentResult {
+    pub fragment_count: usize,
+    pub total_fragment_size_bytes: usize,
+}
+
+impl AddAssign for FlushFragmentResult {
+    fn add_assign(&mut self, rhs: Self) {
+        self.fragment_count += rhs.fragment_count;
+        self.total_fragment_size_bytes += rhs.total_fragment_size_bytes;
+    }
+}
+
+pub struct GenerateEmbeddingsTask {
+    store_client: Arc<dyn StoreClient>,
+    embedding_config: EmbeddingConfig,
+    fragments: Vec<Fragment>,
+    root_node_hash: NodeHash,
+    repo_metadata: RepoMetadata,
+}
+
+pub struct UpdateIntermediateNodesTask {
+    store_client: Arc<dyn StoreClient>,
+    embedding_config: EmbeddingConfig,
+    nodes: Vec<IntermediateNode>,
+}
+
+pub struct SyncMerkleTreeTask {
+    store_client: Arc<dyn StoreClient>,
+    embedding_config: EmbeddingConfig,
+    nodes: Vec<NodeHash>,
+}
+
+pub enum SyncTask {
+    GenerateEmbeddings(GenerateEmbeddingsTask),
+    UpdateIntermediateNodes(UpdateIntermediateNodesTask),
+    SyncMerkleTree(SyncMerkleTreeTask),
+}
+
+#[derive(Debug)]
+pub enum SyncQueueResult {
+    GenerateEmbeddings(HashMap<ContentHash, bool>),
+    UpdateIntermediateNodes(HashMap<NodeHash, bool>),
+    SyncMerkleTree(HashSet<NodeHash>),
+}
+
+impl SyncQueueTaskTrait for SyncTask {
+    type Error = Error;
+    type Result = SyncQueueResult;
+    #[cfg(not(target_arch = "wasm32"))]
+    type Fut = Pin<Box<dyn Future<Output = Result<Self::Result, Self::Error>> + Send>>;
+    #[cfg(target_arch = "wasm32")]
+    type Fut = Pin<Box<dyn Future<Output = Result<Self::Result, Self::Error>>>>;
+    fn run(&mut self) -> Self::Fut {
+        match self {
+            SyncTask::GenerateEmbeddings(task) => {
+                let store_client = task.store_client.clone();
+                let embedding_config = task.embedding_config;
+                let fragments = task.fragments.clone();
+                let root_node_hash = task.root_node_hash.clone();
+                let repo_metadata = task.repo_metadata.clone();
+                Box::pin(async move {
+                    store_client
+                        .generate_embeddings(
+                            embedding_config,
+                            fragments,
+                            root_node_hash,
+                            repo_metadata,
+                        )
+                        .await
+                        .map(SyncQueueResult::GenerateEmbeddings)
+                })
+            }
+            SyncTask::SyncMerkleTree(task) => {
+                let store_client = task.store_client.clone();
+                let embedding_config = task.embedding_config;
+                let nodes = task.nodes.clone();
+                Box::pin(async move {
+                    store_client
+                        .sync_merkle_tree(nodes, embedding_config)
+                        .await
+                        .map(SyncQueueResult::SyncMerkleTree)
+                })
+            }
+            SyncTask::UpdateIntermediateNodes(task) => {
+                let store_client = task.store_client.clone();
+                let embedding_config = task.embedding_config;
+                let nodes = task.nodes.clone();
+                Box::pin(async move {
+                    store_client
+                        .update_intermediate_nodes(embedding_config, nodes)
+                        .await
+                        .map(SyncQueueResult::UpdateIntermediateNodes)
+                })
+            }
+        }
+    }
+}
+
+/// A sync client that is used to update the server merkle tree state so it is up-to-date with the client
+///
+/// A sync is broken down into two steps:
+/// 1) We need to walk the tree to find a list of nodes that need to be updated. This could be done by either
+///    a full scan of the tree or walking the tree from bottom up with a known changed leaf node.
+/// 2) With the list of nodes pending sync known, we could then generate embeddings for the leaf nodes and update
+///    intermediate nodes.
+pub(super) struct CodebaseIndexSyncOperation<'a> {
+    /// A list of dirty nodes that need to be synced with the server. Note that the nodes _MUST_ be ordered from children
+    /// to parent as we cannot sync parents that don't have their children synced.
+    nodes_pending_sync: Vec<NodeLens<'a>>,
+    store_client: Arc<dyn StoreClient>,
+    embedding_config: EmbeddingConfig,
+    sync_queue: SyncQueue<SyncTask>,
+    embedding_generation_batch_size: usize,
+}
+
+impl<'a> CodebaseIndexSyncOperation<'a> {
+    /// Perform a full sync of the merkle tree with the server. This guarantees we will add all inconsistent nodes
+    /// to the nodes_pending_sync list.
+    pub async fn full_sync(
+        tree: &'a MerkleTree,
+        store_client: Arc<dyn StoreClient>,
+        sync_queue: SyncQueue<SyncTask>,
+        sync_progress_tx: async_channel::Sender<SyncProgress>,
+        embedding_generation_batch_size: usize,
+    ) -> Result<(Self, CodebaseContextConfig), SyncOperationError> {
+        let config = store_client
+            .codebase_context_config()
+            .await
+            .map_err(SyncOperationError::ServerSyncError)?;
+
+        let mut operation = Self {
+            nodes_pending_sync: Vec::new(),
+            store_client,
+            embedding_config: config.embedding_config,
+            sync_queue,
+            embedding_generation_batch_size: embedding_generation_batch_size
+                .max(MIN_UPDATE_NODE_BATCH_SIZE),
+        };
+
+        let root_node = tree.root_node();
+        let mut nodes_pending_check = vec![root_node];
+
+        loop {
+            nodes_pending_check = operation
+                .check_if_nodes_synced(&nodes_pending_check, sync_progress_tx.clone())
+                .await?;
+
+            if nodes_pending_check.is_empty() {
+                break;
+            }
+        }
+
+        // We need to reverse the node orders here so we go from children -> parent;
+        operation.nodes_pending_sync.reverse();
+        Ok((operation, config))
+    }
+
+    /// Perform an incremental sync for a set of updated nodes.
+    /// This is used when we know exactly which nodes were modified through file system events.
+    ///
+    /// Returns the operation without flushing the nodes. The caller must call flush_nodes_pending_sync.
+    pub async fn incremental_sync(
+        updated_nodes: Vec<NodeLens<'a>>,
+        store_client: Arc<dyn StoreClient>,
+        sync_queue: SyncQueue<SyncTask>,
+        embedding_config: EmbeddingConfig,
+        sync_progress_tx: async_channel::Sender<SyncProgress>,
+        embedding_generation_batch_size: usize,
+    ) -> Result<Self, SyncOperationError> {
+        if updated_nodes.is_empty() {
+            log::info!("No nodes to sync incrementally");
+        } else {
+            log::debug!(
+                "Starting incremental sync preparation for {} updated nodes",
+                updated_nodes.len()
+            );
+        }
+
+        let mut operation = Self {
+            nodes_pending_sync: Vec::new(),
+            store_client,
+            embedding_config,
+            sync_queue,
+            embedding_generation_batch_size: embedding_generation_batch_size
+                .max(MIN_UPDATE_NODE_BATCH_SIZE),
+        };
+
+        // We need to check if nodes are synced in incremental sync since the to-be-update
+        // nodes could already exist on the server (e.g. switching between git branches).
+        operation
+            .check_if_nodes_synced(&updated_nodes, sync_progress_tx)
+            .await?;
+        Ok(operation)
+    }
+
+    pub async fn flush_nodes_pending_sync(
+        mut self,
+        repo_metadata: &RepoMetadata,
+        root_node_hash: NodeHash,
+        mapping_updates: &LeafToFragmentMetadataMapping,
+        sync_progress_tx: async_channel::Sender<SyncProgress>,
+    ) -> Result<FlushFragmentResult, SyncOperationError> {
+        let mut leaves = Vec::new();
+        let mut intermediate_nodes = Vec::new();
+
+        for node in mem::take(&mut self.nodes_pending_sync) {
+            if node.is_leaf() {
+                leaves.push(node);
+            } else {
+                intermediate_nodes.push(node);
+            }
+        }
+
+        let mut failed_to_sync_nodes: HashSet<NodeHash> = HashSet::new();
+        let mut files_need_resync = ChangedFiles::default();
+        let mut total_fragment_count = 0;
+        let mut total_fragment_size_bytes = 0;
+
+        let total_nodes_to_sync = leaves.len() + intermediate_nodes.len();
+        let mut completed_nodes = 0;
+
+        let leaf_batches = batch_leaves_by_size(
+            &leaves,
+            mapping_updates,
+            self.embedding_generation_batch_size,
+            MAX_BATCH_CONTENT_BYTES,
+        )
+        .map_err(SyncOperationError::Other)?;
+
+        for chunk in &leaf_batches {
+            let mut fragment_metadatas = HashMap::new();
+
+            for node in chunk {
+                let content_hash = node.content_hash().expect("Node should be leaf");
+                let metadatas = mapping_updates
+                    .get(content_hash.as_ref())
+                    .ok_or(anyhow!("Couldn't find metadata for hash"))?;
+
+                for metadata in metadatas {
+                    fragment_metadatas.insert(content_hash.clone(), metadata.clone());
+                }
+            }
+
+            let fragment_metadata_clone = fragment_metadatas.clone();
+            let res = build_fragments_from_metadata(fragment_metadata_clone.into_iter()).await;
+
+            if !res.fail_to_read.is_empty() {
+                let failed_node_count = res.fail_to_read.len();
+                failed_to_sync_nodes.extend(
+                    res.fail_to_read
+                        .into_iter()
+                        .map(|content_hash| content_hash.into()),
+                );
+                files_need_resync.add_paths(res.fail_to_read_path).await;
+                log::warn!("Failed to read {failed_node_count} fragments from disk");
+            }
+
+            // Add retry for embedding generation
+            let fragments_to_sync = res.successfully_read;
+
+            // Track fragment metrics
+            total_fragment_count += fragments_to_sync.len();
+            total_fragment_size_bytes += fragments_to_sync
+                .iter()
+                .map(|f| f.content.len())
+                .sum::<usize>();
+
+            let rx = self
+                .sync_queue
+                .enqueue_with_result(
+                    SyncTask::GenerateEmbeddings(GenerateEmbeddingsTask {
+                        store_client: self.store_client.clone(),
+                        embedding_config: self.embedding_config,
+                        fragments: fragments_to_sync,
+                        root_node_hash: root_node_hash.clone(),
+                        repo_metadata: repo_metadata.clone(),
+                    }),
+                    None,
+                    "generate_embeddings".to_string(),
+                )
+                .await;
+
+            let Ok(task_result) = rx.await else {
+                return Err(SyncOperationError::Other(anyhow::anyhow!(
+                    "Sync queue task cancelled"
+                )));
+            };
+
+            let res = match task_result.inspect(|res| {
+                if let SyncQueueResult::GenerateEmbeddings(res) = res {
+                    let failed_fragments = res
+                        .iter()
+                        .filter_map(|(hash, &success)| {
+                            if !success {
+                                failed_to_sync_nodes.insert(hash.into());
+                                fragment_metadatas.remove(hash)
+                            } else {
+                                None
+                            }
+                        })
+                        .collect_vec();
+                    if !failed_fragments.is_empty() {
+                        log::warn!(
+                            "Failed to generate embeddings for some hashes:\n{failed_fragments:#?}"
+                        );
+                    }
+                }
+            }) {
+                Ok(res) => res,
+                Err(err) => {
+                    log::error!("Failed to generate embeddings: {err:?}");
+                    if files_need_resync.is_empty() {
+                        return Err(SyncOperationError::ServerSyncError(err));
+                    } else {
+                        return Err(SyncOperationError::ReadFragmentError(files_need_resync));
+                    }
+                }
+            };
+
+            completed_nodes += chunk.len();
+            let _ = sync_progress_tx.try_send(SyncProgress::Syncing {
+                completed_nodes: completed_nodes.saturating_sub(failed_to_sync_nodes.len()),
+                total_nodes: total_nodes_to_sync,
+            });
+
+            log::debug!("Generated embedding for the following nodes: {res:?}");
+        }
+
+        for chunk in intermediate_nodes.chunks(self.embedding_generation_batch_size) {
+            let nodes_to_sync: Vec<IntermediateNode> = chunk
+                .iter()
+                .filter_map(|node| {
+                    let outdated = node
+                        .children()
+                        .any(|child| failed_to_sync_nodes.contains(&child.hash()));
+                    if outdated {
+                        failed_to_sync_nodes.insert(node.hash());
+                        None
+                    } else {
+                        Some(IntermediateNode {
+                            hash: node.hash(),
+                            children: node.children().map(|child| child.hash()).collect(),
+                        })
+                    }
+                })
+                .collect();
+
+            let rx = self
+                .sync_queue
+                .enqueue_with_result(
+                    SyncTask::UpdateIntermediateNodes(UpdateIntermediateNodesTask {
+                        store_client: self.store_client.clone(),
+                        embedding_config: self.embedding_config,
+                        nodes: nodes_to_sync.clone(),
+                    }),
+                    None,
+                    "update intermediate nodes".to_string(),
+                )
+                .await;
+
+            let Ok(task_result) = rx.await else {
+                return Err(SyncOperationError::Other(anyhow::anyhow!(
+                    "Sync queue task cancelled"
+                )));
+            };
+
+            let update_result = match task_result.inspect(|res| {
+                if let SyncQueueResult::UpdateIntermediateNodes(res) = res {
+                    let failed_nodes = res
+                        .iter()
+                        .filter_map(|(hash, success)| {
+                            if !success {
+                                failed_to_sync_nodes.insert(hash.clone());
+                                Some(hash.clone())
+                            } else {
+                                None
+                            }
+                        })
+                        .collect_vec();
+
+                    if !failed_nodes.is_empty() {
+                        log::warn!("Failed to sync some intermediate nodes");
+                    }
+                }
+            }) {
+                Ok(res) => res,
+                Err(err) => {
+                    log::error!("Failed to sync intermediate node: {err:?}");
+                    if files_need_resync.is_empty() {
+                        return Err(SyncOperationError::ServerSyncError(err));
+                    } else {
+                        return Err(SyncOperationError::ReadFragmentError(files_need_resync));
+                    }
+                }
+            };
+
+            completed_nodes += chunk.len();
+            let _ = sync_progress_tx.try_send(SyncProgress::Syncing {
+                completed_nodes: completed_nodes.saturating_sub(failed_to_sync_nodes.len()),
+                total_nodes: total_nodes_to_sync,
+            });
+
+            log::debug!("Updated the following nodes: {update_result:?}");
+        }
+
+        if !failed_to_sync_nodes.is_empty() {
+            log::warn!(
+                "Failed to sync {} nodes to the server for root {}",
+                failed_to_sync_nodes.len(),
+                root_node_hash
+            );
+
+            if files_need_resync.is_empty() {
+                return Err(SyncOperationError::ServerSyncError(Error::Other(
+                    anyhow::anyhow!("Failed to sync some nodes to the server"),
+                )));
+            } else {
+                return Err(SyncOperationError::ReadFragmentError(files_need_resync));
+            }
+        }
+
+        log::info!("Successfully flushed nodes pending sync for root {root_node_hash}");
+        Ok(FlushFragmentResult {
+            fragment_count: total_fragment_count,
+            total_fragment_size_bytes,
+        })
+    }
+
+    pub fn total_pending_nodes_count(&self) -> usize {
+        self.nodes_pending_sync.len()
+    }
+
+    async fn check_if_nodes_synced(
+        &mut self,
+        nodes_to_check: &[NodeLens<'a>],
+        sync_progress_tx: async_channel::Sender<SyncProgress>,
+    ) -> Result<Vec<NodeLens<'a>>, SyncOperationError> {
+        let chunks = nodes_to_check.chunks(SYNC_NODE_BATCH_SIZE);
+        let mut res = Vec::new();
+
+        for chunk in chunks {
+            let mut node_hashes = Vec::new();
+
+            for node in chunk {
+                node_hashes.push(node.hash());
+            }
+
+            let rx = self
+                .sync_queue
+                .enqueue_with_result(
+                    SyncTask::SyncMerkleTree(SyncMerkleTreeTask {
+                        store_client: self.store_client.clone(),
+                        embedding_config: self.embedding_config,
+                        nodes: node_hashes.clone(),
+                    }),
+                    None,
+                    "update intermediate nodes".to_string(),
+                )
+                .await;
+
+            let nodes_need_sync = match rx.await {
+                Ok(Ok(SyncQueueResult::SyncMerkleTree(res))) => res,
+                Ok(Ok(_)) => {
+                    return Err(SyncOperationError::Other(anyhow::anyhow!(
+                        "Shouldn't receive other task result in channel"
+                    )))
+                }
+                Ok(Err(e)) => return Err(SyncOperationError::ServerSyncError(e)),
+                Err(_) => {
+                    return Err(SyncOperationError::Other(anyhow::anyhow!(
+                        "Sync queue task cancelled"
+                    )))
+                }
+            };
+
+            // Iterate over the nodes that need to be synced and add their children to the next queue.
+            for node in chunk {
+                if nodes_need_sync.contains(&node.hash()) {
+                    self.nodes_pending_sync.push(*node);
+                    res.extend(node.children());
+                }
+            }
+        }
+
+        let _ = sync_progress_tx.try_send(SyncProgress::Discovering {
+            total_nodes: self.nodes_pending_sync.len(),
+        });
+
+        Ok(res)
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub(super) enum SyncOperationError {
+    #[error("Error reading some fragments {0:#?}")]
+    ReadFragmentError(ChangedFiles),
+    #[error("Error syncing nodes with server {0:#}")]
+    ServerSyncError(Error),
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+}
+
+/// Partitions `leaves` into batches where each batch contains at most `max_count` leaves
+/// AND at most `max_bytes` of estimated content.
+///
+/// Content size for each leaf is estimated from the `byte_range` in its `FragmentMetadata`,
+/// which is available without reading from disk.
+fn batch_leaves_by_size<'a>(
+    leaves: &[NodeLens<'a>],
+    mapping_updates: &LeafToFragmentMetadataMapping,
+    max_count: usize,
+    max_bytes: usize,
+) -> Result<Vec<Vec<NodeLens<'a>>>> {
+    if leaves.is_empty() {
+        return Ok(vec![]);
+    }
+
+    let mut batches = Vec::new();
+    let mut current_batch = Vec::new();
+    let mut current_bytes: usize = 0;
+
+    for leaf in leaves {
+        let content_hash = leaf.content_hash().expect("Node should be leaf");
+        let metadatas = mapping_updates
+            .get(content_hash.as_ref())
+            .ok_or_else(|| anyhow!("Couldn't find metadata for hash {content_hash:?}"))?;
+
+        let leaf_bytes: usize = metadatas.iter().map(|m| m.content_byte_size()).sum();
+
+        // If the current batch is non-empty and adding this leaf would exceed either limit,
+        // finalize the current batch and start a new one.
+        if !current_batch.is_empty()
+            && (current_batch.len() >= max_count || current_bytes + leaf_bytes > max_bytes)
+        {
+            batches.push(std::mem::take(&mut current_batch));
+            current_bytes = 0;
+        }
+
+        current_batch.push(*leaf);
+        current_bytes += leaf_bytes;
+    }
+
+    if !current_batch.is_empty() {
+        batches.push(current_batch);
+    }
+
+    Ok(batches)
+}
+
+impl IsTransientError for Error {
+    fn is_transient(&self) -> bool {
+        // TODO: match on the error type and only return true of actual transient error.
+        true
+    }
+}
+
+impl From<SyncOperationError> for Error {
+    fn from(error: SyncOperationError) -> Self {
+        match error {
+            SyncOperationError::Other(e) => Self::Other(e),
+            SyncOperationError::ServerSyncError(e) => e,
+            SyncOperationError::ReadFragmentError(_) => Self::FileSystemStateChanged,
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "sync_client_tests.rs"]
+mod tests;

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

@@ -0,0 +1,112 @@
+use std::collections::HashMap;
+
+use futures::executor::block_on;
+
+use super::batch_leaves_by_size;
+use crate::index::full_source_code_embedding::merkle_tree::{construct_test_merkle_tree, NodeLens};
+
+use virtual_fs::VirtualFS;
+
+/// Collect all leaf nodes from a merkle tree by walking it recursively.
+fn collect_leaves<'a>(node: NodeLens<'a>) -> Vec<NodeLens<'a>> {
+    if node.is_leaf() {
+        return vec![node];
+    }
+    node.children().flat_map(collect_leaves).collect()
+}
+
+#[test]
+fn test_batch_leaves_single_batch_when_under_limits() {
+    VirtualFS::test("batch_single", |dirs, mut sandbox| {
+        let (tree, metadata) = block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+
+        let leaves = collect_leaves(tree.root_node());
+        assert!(!leaves.is_empty(), "Tree should have leaf nodes");
+
+        let batches = batch_leaves_by_size(&leaves, metadata.mapping(), 1000, 10_000_000).unwrap();
+
+        assert_eq!(batches.len(), 1, "All leaves should fit in a single batch");
+        assert_eq!(
+            batches[0].len(),
+            leaves.len(),
+            "The single batch should contain all leaves"
+        );
+    });
+}
+
+#[test]
+fn test_batch_leaves_splits_on_count_limit() {
+    VirtualFS::test("batch_count", |dirs, mut sandbox| {
+        let (tree, metadata) = block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+
+        let leaves = collect_leaves(tree.root_node());
+        let leaf_count = leaves.len();
+        assert!(leaf_count >= 2, "Need at least 2 leaves for this test");
+
+        let batches = batch_leaves_by_size(&leaves, metadata.mapping(), 1, 10_000_000).unwrap();
+
+        assert_eq!(
+            batches.len(),
+            leaf_count,
+            "Each leaf should be in its own batch when max_count=1"
+        );
+        for batch in &batches {
+            assert_eq!(batch.len(), 1);
+        }
+    });
+}
+
+#[test]
+fn test_batch_leaves_splits_on_byte_limit() {
+    VirtualFS::test("batch_bytes", |dirs, mut sandbox| {
+        let (tree, metadata) = block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+
+        let leaves = collect_leaves(tree.root_node());
+        let leaf_count = leaves.len();
+        assert!(leaf_count >= 2, "Need at least 2 leaves for this test");
+
+        // max_bytes=1 means every leaf exceeds the limit, but progress guarantee
+        // ensures each still gets its own batch.
+        let batches = batch_leaves_by_size(&leaves, metadata.mapping(), 1000, 1).unwrap();
+
+        assert_eq!(
+            batches.len(),
+            leaf_count,
+            "Each leaf should be in its own batch when max_bytes=1 (progress guarantee)"
+        );
+        for batch in &batches {
+            assert_eq!(batch.len(), 1);
+        }
+    });
+}
+
+#[test]
+fn test_batch_leaves_empty_input() {
+    let leaves: Vec<NodeLens<'_>> = vec![];
+    let mapping = HashMap::new();
+
+    let batches = batch_leaves_by_size(&leaves, &mapping, 100, 4_000_000).unwrap();
+    assert!(
+        batches.is_empty(),
+        "Empty input should produce empty output"
+    );
+}
+
+#[test]
+fn test_batch_leaves_missing_metadata_returns_error() {
+    VirtualFS::test("batch_missing", |dirs, mut sandbox| {
+        let (tree, _metadata) = block_on(construct_test_merkle_tree(&dirs, &mut sandbox));
+
+        let leaves = collect_leaves(tree.root_node());
+        assert!(!leaves.is_empty());
+
+        // Pass an empty mapping — every leaf lookup should fail.
+        let empty_mapping = HashMap::new();
+        let result = batch_leaves_by_size(&leaves, &empty_mapping, 1000, 10_000_000);
+
+        assert!(
+            result.is_err(),
+            "Should return error when metadata is missing"
+        );
+    });
+}