Add local project indexing and search guidance

This commit is contained in:
2026-08-30 22:00:27 -05:00
parent 88c1ef9716
commit 1c7d3c175d
39 changed files with 2094 additions and 306 deletions
+264 -77
View File
@@ -2,25 +2,38 @@ use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use ai::index::full_source_code_embedding::manager::{
CodebaseIndexManager, CodebaseIndexManagerEvent,
};
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use ai::index::full_source_code_embedding::RetrievalID;
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
use ai::index::local_project_index::LocalProjectIndexManager;
use ai::index::locations::CodeContextLocation;
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use anyhow::anyhow;
use futures_util::stream::AbortHandle;
use galaxy_core::features::FeatureFlag;
use galaxyui::{AppContext, Entity, ModelContext, ModelHandle, SingletonEntity};
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use instant::Instant;
#[cfg(not(target_family = "wasm"))]
use crate::ai::agent::SearchCodebaseFailureReason;
use crate::ai::agent::{AIAgentActionId, SearchCodebaseResult};
use crate::ai::blocklist::SessionContext;
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use crate::ai::get_relevant_files::api::{FileContext as FileContextRequest, GetRelevantFiles};
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
use crate::ai::get_relevant_files::controller::GetRelevantFilesError::LocalIndexUnavailable;
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use crate::ai::outline::{OutlineStatus, RepoOutlines};
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use crate::report_error;
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
use crate::server::server_api::{AIApiError, ServerApiProvider};
use crate::{report_error, send_telemetry_from_ctx, TelemetryEvent};
use crate::{send_telemetry_from_ctx, TelemetryEvent};
#[cfg_attr(not(target_family = "wasm"), path = "remote_search/native.rs")]
#[cfg_attr(target_family = "wasm", path = "remote_search/wasm.rs")]
mod remote_search;
@@ -60,6 +73,20 @@ pub enum GetRelevantFilesRequestTarget {
requested_codebase_path: Option<String>,
},
}
/// Lifecycle state for the two-stage `get_files` relevant-file query.
///
/// The first stage discovers paths through `GetRelevantFilesController`; the second stage reads
/// current file contents in `GetFilesExecutor`. Keeping this state separate from the controller's
/// cancellation handles preserves the existing action protocol for both local and legacy search.
#[derive(Debug, Clone)]
pub enum GetRelevantFilesStatus {
Pending { root_repo_path: PathBuf },
InFlight { root_repo_path: PathBuf },
Success { file_paths: Vec<PathBuf> },
Failed { message: String },
}
#[derive(Debug, thiserror::Error)]
pub enum GetRelevantFilesError {
#[error("Repo outline is still being computed.")]
@@ -68,15 +95,18 @@ pub enum GetRelevantFilesError {
CreateFailed,
#[error("Failed to create outline.")]
Missing,
#[error("Local project index is unavailable: {0}")]
LocalIndexUnavailable(String),
}
/// This enum allows us to use both the existing structure for outline-based indexing
/// and the new full source code indexing manager/model.
enum RequestHandle {
/// Used with outline-based indexing.
/// Used with outline-based indexing and remote search.
AbortHandle(AbortHandle),
/// Used with full source code indexing.
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
RetrievalID {
repo_path: PathBuf,
retrieval_id: RetrievalID,
@@ -88,6 +118,7 @@ impl RequestHandle {
fn abort(&mut self, ctx: &mut AppContext) {
match self {
RequestHandle::AbortHandle(abort_handle) => abort_handle.abort(),
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
RequestHandle::RetrievalID {
repo_path,
retrieval_id,
@@ -112,15 +143,63 @@ pub struct GetRelevantFilesController {
/// This allows several SearchCodebase actions to be active at once without newer requests
/// cancelling unrelated older ones.
pending_requests: std::collections::HashMap<AIAgentActionId, RequestHandle>,
/// Status for the two-stage `get_files` relevant-file query protocol.
relevant_file_statuses: std::collections::HashMap<AIAgentActionId, GetRelevantFilesStatus>,
}
impl GetRelevantFilesController {
pub fn new(ctx: &mut ModelContext<Self>) -> Self {
let codebase_manager = CodebaseIndexManager::handle(ctx);
ctx.subscribe_to_model(&codebase_manager, Self::handle_codebase_manager_event);
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
{
let codebase_manager = CodebaseIndexManager::handle(ctx);
ctx.subscribe_to_model(&codebase_manager, Self::handle_codebase_manager_event);
}
Self::default()
}
fn root_for_status(&self, directory: &Path, app: &AppContext) -> PathBuf {
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
{
app.try_get_singleton_model_as_ref::<LocalProjectIndexManager>()
.and_then(|manager| manager.root_for_path(directory))
.unwrap_or_else(|| directory.to_path_buf())
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
{
self.root_directory_for_search(directory, app)
.unwrap_or_else(|| directory.to_path_buf())
}
}
/// Queues the first stage of a relevant-file query for `GetFilesExecutor`.
pub fn queue_request(
&mut self,
action_id: AIAgentActionId,
directory: &Path,
app: &AppContext,
) {
let root_repo_path = self.root_for_status(directory, app);
self.relevant_file_statuses.insert(
action_id,
GetRelevantFilesStatus::Pending { root_repo_path },
);
}
/// Returns the current status of a relevant-file query.
pub fn status(&self, action_id: &AIAgentActionId) -> Option<&GetRelevantFilesStatus> {
self.relevant_file_statuses.get(action_id)
}
fn set_relevant_file_status(
&mut self,
action_id: &AIAgentActionId,
status: GetRelevantFilesStatus,
) {
self.relevant_file_statuses
.insert(action_id.clone(), status);
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
fn pending_request_details_for_retrieval_id(
&self,
pending_retrieval_id: &RetrievalID,
@@ -140,6 +219,7 @@ impl GetRelevantFilesController {
})
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
fn handle_codebase_manager_event(
&mut self,
_: ModelHandle<CodebaseIndexManager>,
@@ -213,7 +293,38 @@ impl GetRelevantFilesController {
self.cancel_request_for_action(&action_id, ctx);
match target {
GetRelevantFilesRequestTarget::Local { directory } => {
self.send_local_request(&directory, query, partial_path_segments, action_id, ctx)
let root_repo_path = self
.relevant_file_statuses
.get(&action_id)
.and_then(|status| match status {
GetRelevantFilesStatus::Pending { root_repo_path }
| GetRelevantFilesStatus::InFlight { root_repo_path } => {
Some(root_repo_path.clone())
}
GetRelevantFilesStatus::Success { .. }
| GetRelevantFilesStatus::Failed { .. } => None,
})
.unwrap_or_else(|| self.root_for_status(&directory, ctx));
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::InFlight { root_repo_path },
);
let result = self.send_local_request(
&directory,
query,
partial_path_segments,
action_id.clone(),
ctx,
);
if let Err(error) = &result {
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::Failed {
message: error.to_string(),
},
);
}
result
}
GetRelevantFilesRequestTarget::Remote {
session_context,
@@ -237,8 +348,45 @@ impl GetRelevantFilesController {
action_id: AIAgentActionId,
ctx: &mut ModelContext<Self>,
) -> Result<(), GetRelevantFilesError> {
const MINIMUM_FILE_COUNT_FOR_API_CALL: usize = 2;
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
{
let Some(local_index) =
ctx.try_get_singleton_model_as_ref::<LocalProjectIndexManager>()
else {
return Err(LocalIndexUnavailable(
"The local project index is not available in this session.".to_string(),
));
};
if local_index.is_searchable_for_path(directory) {
let locations: HashSet<CodeContextLocation> = local_index
.search(directory, &query, partial_path_segments.map(Vec::as_slice))
.map_err(|error| LocalIndexUnavailable(error.to_string()))?
.into_iter()
.map(|hit| CodeContextLocation::WholeFile(hit.path))
.collect();
let file_paths = locations
.iter()
.map(|location| location.path().clone())
.collect();
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::Success { file_paths },
);
ctx.emit(GetRelevantFilesControllerEvent::Success {
action_id,
result: GetRelevantFilesControllerResult::Locations(Arc::new(locations)),
});
return Ok(());
}
return Err(LocalIndexUnavailable(
"The local project index is not ready. Run /init or /index first.".to_string(),
));
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
let _ = (directory, query, partial_path_segments, action_id, ctx);
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
if FeatureFlag::FullSourceCodeEmbedding.is_enabled() {
let codebase_mgr = CodebaseIndexManager::handle(ctx);
if let Some(base_path) = codebase_mgr.as_ref(ctx).root_path_for_codebase(directory) {
@@ -268,79 +416,93 @@ impl GetRelevantFilesController {
}
}
match RepoOutlines::as_ref(ctx).get_outline(directory) {
Some((OutlineStatus::Complete(outline), base_path)) => {
let server_api = ServerApiProvider::as_ref(ctx).get();
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
{
match RepoOutlines::as_ref(ctx).get_outline(directory) {
Some((OutlineStatus::Complete(outline), base_path)) => {
let server_api = ServerApiProvider::as_ref(ctx).get();
let file_outlines = outline.to_file_symbols(partial_path_segments);
if file_outlines.len() < MINIMUM_FILE_COUNT_FOR_API_CALL {
ctx.emit(GetRelevantFilesControllerEvent::Success {
action_id,
result: GetRelevantFilesControllerResult::Locations(Arc::new(
let file_outlines = outline.to_file_symbols(partial_path_segments);
if file_outlines.len() < MINIMUM_FILE_COUNT_FOR_API_CALL {
let locations: Arc<HashSet<CodeContextLocation>> = Arc::new(
file_outlines
.into_iter()
.map(|file| {
CodeContextLocation::WholeFile(PathBuf::from(file.path))
})
.collect(),
)),
});
} else {
let outline_request = GetRelevantFiles {
query,
files: file_outlines
.into_iter()
.map(|outline| FileContextRequest {
path: outline.path,
symbols: outline.symbols,
})
.collect(),
};
let action_id_clone = action_id.clone();
let request_abort_handle = ctx
.spawn(
async move {
let response =
server_api.get_relevant_files(&outline_request).await?;
Ok(Arc::new(
response
.relevant_file_paths
.into_iter()
.filter_map(|path| {
let file_path = base_path.join(path);
// Validate the returned file paths.
if file_path.exists() {
Some(CodeContextLocation::WholeFile(file_path))
} else {
None
}
})
.collect(),
))
},
move |me,
relevant_file_paths: Result<
Arc<HashSet<CodeContextLocation>>,
AIApiError,
>,
ctx| {
me.handle_relevant_file_paths_result(
relevant_file_paths.map_err(|e| anyhow!(e)),
action_id_clone,
ctx,
)
},
)
.abort_handle();
self.pending_requests
.insert(action_id, RequestHandle::AbortHandle(request_abort_handle));
);
let file_paths = locations
.iter()
.map(|location| location.path().clone())
.collect();
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::Success { file_paths },
);
ctx.emit(GetRelevantFilesControllerEvent::Success {
action_id,
result: GetRelevantFilesControllerResult::Locations(locations),
});
} else {
let outline_request = GetRelevantFiles {
query,
files: file_outlines
.into_iter()
.map(|outline| FileContextRequest {
path: outline.path,
symbols: outline.symbols,
})
.collect(),
};
let action_id_clone = action_id.clone();
let request_abort_handle = ctx
.spawn(
async move {
let response =
server_api.get_relevant_files(&outline_request).await?;
Ok(Arc::new(
response
.relevant_file_paths
.into_iter()
.filter_map(|path| {
let file_path = base_path.join(path);
// Validate the returned file paths.
if file_path.exists() {
Some(CodeContextLocation::WholeFile(file_path))
} else {
None
}
})
.collect(),
))
},
move |me,
relevant_file_paths: Result<
Arc<HashSet<CodeContextLocation>>,
AIApiError,
>,
ctx| {
me.handle_relevant_file_paths_result(
relevant_file_paths.map_err(|e| anyhow!(e)),
action_id_clone,
ctx,
)
},
)
.abort_handle();
self.pending_requests
.insert(action_id, RequestHandle::AbortHandle(request_abort_handle));
}
Ok(())
}
Ok(())
Some((OutlineStatus::Pending, _)) => Err(GetRelevantFilesError::Pending),
Some((OutlineStatus::Failed, _)) => Err(GetRelevantFilesError::CreateFailed),
None => Err(GetRelevantFilesError::Missing),
}
Some((OutlineStatus::Pending, _)) => Err(GetRelevantFilesError::Pending),
Some((OutlineStatus::Failed, _)) => Err(GetRelevantFilesError::CreateFailed),
None => Err(GetRelevantFilesError::Missing),
}
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
unreachable!("native local search returns before the legacy outline fallback")
}
fn send_remote_request(
@@ -375,6 +537,7 @@ impl GetRelevantFilesController {
Ok(())
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
fn handle_relevant_file_paths_result(
&mut self,
relevant_file_locations: anyhow::Result<Arc<HashSet<CodeContextLocation>>>,
@@ -386,12 +549,27 @@ impl GetRelevantFilesController {
}
match relevant_file_locations {
Ok(relevant_file_locations) => {
let file_paths = relevant_file_locations
.iter()
.map(|location| location.path().clone())
.collect();
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::Success { file_paths },
);
ctx.emit(GetRelevantFilesControllerEvent::Success {
action_id,
result: GetRelevantFilesControllerResult::Locations(relevant_file_locations),
});
}
Err(e) => {
let message = e.to_string();
self.set_relevant_file_status(
&action_id,
GetRelevantFilesStatus::Failed {
message: message.clone(),
},
);
report_error!(anyhow!(e).context("get_relevant_files failed"));
ctx.emit(GetRelevantFilesControllerEvent::Error { action_id });
}
@@ -421,15 +599,24 @@ impl GetRelevantFilesController {
/// Returns the path to the root directory for a codebase search where pwd is `directory`.
pub fn root_directory_for_search(&self, directory: &Path, app: &AppContext) -> Option<PathBuf> {
let mut start = None;
if FeatureFlag::FullSourceCodeEmbedding.is_enabled() {
start = CodebaseIndexManager::as_ref(app).root_path_for_codebase(directory);
#[cfg(all(feature = "local_fs", not(target_family = "wasm")))]
{
app.try_get_singleton_model_as_ref::<LocalProjectIndexManager>()
.and_then(|manager| manager.root_for_path(directory))
}
#[cfg(any(not(feature = "local_fs"), target_family = "wasm"))]
{
let mut start = None;
if FeatureFlag::FullSourceCodeEmbedding.is_enabled() {
start = CodebaseIndexManager::as_ref(app).root_path_for_codebase(directory);
}
start.or_else(|| {
RepoOutlines::as_ref(app)
.get_outline(directory)
.map(|(_, root)| root)
})
}
start.or_else(|| {
RepoOutlines::as_ref(app)
.get_outline(directory)
.map(|(_, root)| root)
})
}
pub fn root_directory_for_remote_search(