public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reduce indexing overhead and repair agent search and child recovery

Browse Source
This commit is contained in:
rkw6086
2026-09-14 16:25:49 -05:00
parent 81b4c78b61
commit 506dc5b225
38 changed files with 1941 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model/execute/search_codebase.rs
+50 -16
View File
@@ -27,12 +27,21 @@ use crate::features::FeatureFlag;
use crate::terminal::model::session::active_session::ActiveSession;
use crate::{send_telemetry_from_ctx, TelemetryEvent};
#[path = "search_codebase_context.rs"]
mod context;
use context::{bound_search_result, MAX_SEARCH_FILES, MAX_SEARCH_SCAN_BYTES};
struct PendingSearch {
sender: oneshot::Sender<SearchCodebaseResult>,
query: String,
}
pub struct SearchCodebaseExecutor {
active_session: ModelHandle<ActiveSession>,
get_relevant_files_controller: ModelHandle<GetRelevantFilesController>,
/// Per-action response channels for searches that are still waiting on
/// `GetRelevantFilesController`.
active_searches: HashMap<AIAgentActionId, oneshot::Sender<SearchCodebaseResult>>,
active_searches: HashMap<AIAgentActionId, PendingSearch>,
/// Cached repo roots derived during preprocessing so permission checks and execution can agree
/// on which repository the action actually targets.
root_repo_paths: HashMap<AIAgentActionId, PathBuf>,
@@ -59,9 +68,7 @@ impl SearchCodebaseExecutor {
let Some(result_tx) = me.active_searches.remove(action_id) else {
return;
};
if let Err(e) = result_tx.send(result.clone()) {
log::warn!("Failed to send search codebase results to receiver {e:?}.");
}
Self::send_bounded_result(result_tx, result.clone(), ctx);
}
GetRelevantFilesControllerEvent::Success {
result: GetRelevantFilesControllerResult::Locations(fragments),
@@ -70,6 +77,8 @@ impl SearchCodebaseExecutor {
let action_id = event.action_id().clone();
let locations = fragments
.iter()
.sorted_by_key(|location| location.path())
.take(MAX_SEARCH_FILES)
.map(|location| location.into())
.collect_vec();
let current_working_directory = me
@@ -84,13 +93,15 @@ impl SearchCodebaseExecutor {
&locations,
current_working_directory,
shell,
None,
None,
Some(MAX_SEARCH_SCAN_BYTES),
Some(MAX_SEARCH_SCAN_BYTES * MAX_SEARCH_FILES),
)
.await
{
Ok(result) => {
if !result.missing_files.is_empty() {
if result.file_contexts.is_empty()
&& !result.missing_files.is_empty()
{
let missing_files = result.missing_files.join(", ");
SearchCodebaseResult::Failed {
message: format!(
@@ -110,15 +121,11 @@ impl SearchCodebaseExecutor {
},
}
},
move |me, result, _| {
move |me, result, ctx| {
let Some(result_tx) = me.active_searches.remove(&action_id) else {
return;
};
if let Err(e) = result_tx.send(result) {
log::warn!(
"Failed to send search codebase results to receiver {e:?}."
);
}
Self::send_bounded_result(result_tx, result, ctx);
},
);
}
@@ -126,7 +133,7 @@ impl SearchCodebaseExecutor {
let Some(result_tx) = me.active_searches.remove(action_id) else {
return;
};
if let Err(e) = result_tx.send(SearchCodebaseResult::Failed {
if let Err(e) = result_tx.sender.send(SearchCodebaseResult::Failed {
message: "The search failed. Try another way to locate the relevant files."
.to_owned(),
reason: SearchCodebaseFailureReason::GetRelevantFilesError,
@@ -146,6 +153,21 @@ impl SearchCodebaseExecutor {
}
}
fn send_bounded_result(
pending: PendingSearch,
result: SearchCodebaseResult,
ctx: &mut ModelContext<Self>,
) {
ctx.spawn(
async move { (pending.sender, bound_search_result(result, &pending.query)) },
|_, (sender, result), _| {
if let Err(error) = sender.send(result) {
log::debug!("Search result receiver was dropped: {error:?}");
}
},
);
}
pub(super) fn should_autoexecute(
&self,
input: ExecuteActionInput,
@@ -245,7 +267,13 @@ impl SearchCodebaseExecutor {
});
let (result_tx, result_rx) = oneshot::channel();
self.active_searches.insert(id.clone(), result_tx);
self.active_searches.insert(
id.clone(),
PendingSearch {
sender: result_tx,
query: query.clone(),
},
);
match self
.get_relevant_files_controller
@@ -358,7 +386,13 @@ impl SearchCodebaseExecutor {
});
let (result_tx, result_rx) = oneshot::channel();
self.active_searches.insert(id.clone(), result_tx);
self.active_searches.insert(
id.clone(),
PendingSearch {
sender: result_tx,
query: query.clone(),
},
);
// Start the actual search.
match self
app/src/ai/blocklist/action_model/execute/search_codebase_context.rs
+121
View File
@@ -0,0 +1,121 @@
//! Bounded discovery excerpts. Explicit `read_files` calls provide full follow-up context.
use std::collections::HashSet;
pub(super) use ai::agent::action_result::MAX_SEARCH_CODEBASE_FILES as MAX_SEARCH_FILES;
use crate::ai::agent::{AnyFileContent, FileContext, SearchCodebaseResult};
pub(super) const MAX_SEARCH_SCAN_BYTES: usize = 256_000;
const MAX_EXCERPT_BYTES: usize = 4_000;
const MAX_EXCERPT_LINES: usize = 40;
pub(super) fn bound_search_result(
result: SearchCodebaseResult,
query: &str,
) -> SearchCodebaseResult {
let SearchCodebaseResult::Success { files } = result else {
return result;
};
let terms: Vec<_> = query
.split(|character: char| !character.is_alphanumeric() && character != '_')
.filter(|term| term.len() >= 3)
.map(str::to_lowercase)
.filter(|term| {
!matches!(
term.as_str(),
"the"
| "and"
| "for"
| "with"
| "from"
| "that"
| "this"
| "find"
| "show"
| "all"
| "code"
| "files"
| "where"
| "what"
| "when"
| "which"
| "into"
)
})
.take(32)
.collect();
let mut seen = HashSet::new();
let files = files
.into_iter()
.filter(|file| seen.insert(file.file_name.clone()))
.filter_map(|file| excerpt(file, &terms))
.take(MAX_SEARCH_FILES)
.collect();
SearchCodebaseResult::Success { files }
}
fn excerpt(mut file: FileContext, terms: &[String]) -> Option<FileContext> {
let AnyFileContent::StringContent(content) = &file.content else {
return None;
};
// Also bound remote payloads before scoring. Never split a UTF-8 code point.
let mut scan_end = content.len().min(MAX_SEARCH_SCAN_BYTES);
while !content.is_char_boundary(scan_end) {
scan_end -= 1;
}
let lines: Vec<_> = content[..scan_end].split_inclusive('\n').collect();
if lines.is_empty() {
return None;
}
let mut best_line = 0;
let mut best_score = 0;
for (index, line) in lines.iter().enumerate() {
let lower = line.to_lowercase();
let score: usize = terms
.iter()
.filter(|term| lower.contains(term.as_str()))
.map(String::len)
.sum();
if score > best_score {
best_score = score;
best_line = index;
}
}
let mut start = best_line;
let mut preceding_bytes = 0;
while start > best_line.saturating_sub(5) {
let length = lines[start - 1].len();
if preceding_bytes + length > MAX_EXCERPT_BYTES / 3 {
break;
}
preceding_bytes += length;
start -= 1;
}
let mut text = String::new();
let mut line_count: usize = 0;
for line in lines.iter().skip(start).take(MAX_EXCERPT_LINES) {
let remaining = MAX_EXCERPT_BYTES - text.len();
if remaining == 0 {
break;
}
let mut end = line.len().min(remaining);
while !line.is_char_boundary(end) {
end -= 1;
}
if end == 0 {
break;
}
text.push_str(&line[..end]);
line_count += 1;
if end < line.len() {
break;
}
}
let first_line = file.line_range.as_ref().map_or(1, |range| range.start) + start;
file.line_range = Some(first_line..first_line + line_count.saturating_sub(1));
file.content = AnyFileContent::StringContent(text);
Some(file)
}
#[cfg(test)]
#[path = "search_codebase_context_tests.rs"]
mod tests;
app/src/ai/blocklist/action_model/execute/search_codebase_context_tests.rs
+68
View File
@@ -0,0 +1,68 @@
use super::*;
fn file(name: &str, text: String) -> FileContext {
FileContext::new(
name.to_string(),
AnyFileContent::StringContent(text),
None,
None,
)
}
#[test]
fn search_returns_bounded_unique_files_with_valid_utf8() {
let files = (0..20)
.flat_map(|index| {
let file = file(&format!("{index}.rs"), "€".repeat(20_000));
[file.clone(), file]
})
.collect();
let SearchCodebaseResult::Success { files } =
bound_search_result(SearchCodebaseResult::Success { files }, "find models")
else {
panic!("expected success");
};
assert_eq!(files.len(), MAX_SEARCH_FILES);
let mut bytes = 0;
for file in files {
let AnyFileContent::StringContent(text) = file.content else {
panic!("expected text");
};
assert!(text.len() <= MAX_EXCERPT_BYTES);
assert_eq!(text.len() % 3, 0);
bytes += text.len();
}
assert!(bytes <= MAX_SEARCH_FILES * MAX_EXCERPT_BYTES);
}
#[test]
fn excerpt_finds_a_later_match_and_preserves_source_line_numbers() {
let text = format!(
"{}fn construct_chat_model() {{}}\n{}",
"// unrelated\n".repeat(200),
"// trailing\n".repeat(100)
);
let mut input = file("models.rs", text);
input.line_range = Some(101..401);
let output = excerpt(input, &["construct_chat_model".into()]).unwrap();
assert_eq!(output.line_range, Some(296..335));
let AnyFileContent::StringContent(text) = output.content else {
panic!("expected text");
};
assert!(text.contains("construct_chat_model"));
assert_eq!(text.lines().count(), MAX_EXCERPT_LINES);
}
#[test]
fn long_preceding_line_does_not_hide_the_relevant_match() {
let input = file(
"models.rs",
format!("{}\nfn target_model() {{}}\n", "x".repeat(10_000)),
);
let output = excerpt(input, &["target_model".into()]).unwrap();
assert_eq!(output.line_range, Some(2..2));
let AnyFileContent::StringContent(text) = output.content else {
panic!("expected text");
};
assert!(text.contains("target_model"));
}
app/src/ai/blocklist/action_model/execute/send_message.rs
+67 -1
View File
@@ -1,3 +1,4 @@
use std::collections::HashMap;
#[cfg(not(target_family = "wasm"))]
use std::time::Duration;
@@ -13,12 +14,13 @@ use warpui::r#async::Timer;
use warpui::{AppContext, Entity, ModelContext, SingletonEntity};
use super::{ActionExecution, AnyActionExecution, ExecuteActionInput, PreprocessActionInput};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent::conversation::{AIConversationId, ConversationStatus};
use crate::ai::agent::{
AIAgentAction, AIAgentActionResultType, AIAgentActionType, SendMessageToAgentResult,
};
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::blocklist::history_model::BlocklistAIHistoryModel;
use crate::ai::blocklist::orchestration_events::OrchestrationEventService;
use crate::ai::blocklist::telemetry::{
BlocklistOrchestrationTelemetryEvent, TeamAgentCommunicationFailedEvent,
TeamAgentCommunicationFailureReason, TeamAgentCommunicationKind,
@@ -32,6 +34,7 @@ const SEND_AGENT_MESSAGE_TIMEOUT: Duration = Duration::from_secs(15);
pub struct SendMessageToAgentExecutor {
ambient_agent_task_id: Option<AmbientAgentTaskId>,
local_recovery_attempts: HashMap<AIConversationId, usize>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -121,6 +124,7 @@ impl SendMessageToAgentExecutor {
pub fn new() -> Self {
Self {
ambient_agent_task_id: None,
local_recovery_attempts: HashMap::new(),
}
}
@@ -159,6 +163,68 @@ impl SendMessageToAgentExecutor {
let subject = subject.clone();
let message_body = message.clone();
// Resolve only an actual child of this sender; local routing never trusts an arbitrary
// conversation ID as authority to resume another task.
let history = BlocklistAIHistoryModel::as_ref(ctx);
let local_children = history.child_conversations_of(conversation_id);
let local_child = if addresses.len() == 1 {
local_children.iter().find(|child| {
!child.is_remote_child()
&& (child.id().to_string() == addresses[0]
|| child.orchestration_agent_id().as_deref() == Some(addresses[0].as_str()))
})
} else {
None
};
if let Some(child) = local_child {
let child_id = child.id();
let error = if message_body.trim().is_empty() {
Some("A corrective instruction or follow-up message is required.")
} else if matches!(child.status(), ConversationStatus::Cancelled) {
Some("The child was cancelled. Ask the user to resume it; do not override cancellation.")
} else if matches!(child.status(), ConversationStatus::Blocked { .. }) {
Some("The child is blocked on permission or user input. Resolve that block rather than restarting it.")
} else if matches!(child.status(), ConversationStatus::Error)
&& self
.local_recovery_attempts
.get(&child_id)
.copied()
.unwrap_or(0)
>= 2
{
Some("This child has exhausted two orchestrator recovery attempts. Report the blocker to the user.")
} else {
None
};
if let Some(error) = error {
return ActionExecution::<()>::Sync(AIAgentActionResultType::SendMessageToAgent(
SendMessageToAgentResult::Error(error.to_string()),
))
.into();
}
if matches!(child.status(), ConversationStatus::Error) {
*self.local_recovery_attempts.entry(child_id).or_default() += 1;
}
let message_id = OrchestrationEventService::handle(ctx).update(ctx, |service, ctx| {
service.send_local_follow_up(conversation_id, child_id, message_body, ctx)
});
return ActionExecution::<()>::Sync(AIAgentActionResultType::SendMessageToAgent(
SendMessageToAgentResult::Success { message_id },
))
.into();
}
if history
.conversation(&conversation_id)
.is_some_and(|parent| parent.run_id().is_none())
|| addresses.iter().any(|address| {
history.all_live_conversations().iter().any(|(_, target)| {
!target.is_remote_child() && target.id().to_string() == *address
})
})
{
return ActionExecution::<()>::Sync(AIAgentActionResultType::SendMessageToAgent(SendMessageToAgentResult::Error("The target must be one existing child of this orchestrator. Use its agent_id from run_agents; do not message unrelated conversations.".to_string()))).into();
}
let (sender_run_id, task_id, task_resolution) =
sender_run_id_and_task_id_for_send(conversation_id, self.ambient_agent_task_id, ctx);
let log_addresses = addresses.clone();
app/src/ai/blocklist/action_model/execute/start_agent.rs
+5
View File
@@ -313,6 +313,11 @@ impl StartAgentExecutor {
let Some(pending) = self.pending.remove(&request_id) else {
return;
};
let error_msg = if matches!(pending.wait_policy, StartAgentWaitPolicy::Completion) {
format!("Child agent {child_conversation_id} failed: {error_msg}. For a recoverable error (never a user cancellation), use send_message_to_agent with this agent_id and a corrective instruction to continue the same child; do not create a duplicate.")
} else {
error_msg
};
let _ = pending.sender.try_send(StartAgentOutcome::Error(error_msg));
// Only startup acknowledgements may clean up a conversation that never
// initialized. Direct-provider completion waits preserve the terminal