assistant: Update SlashCommand trait with streaming return type (#19652)

This PR updates the `SlashCommand` trait to use a streaming return type.

This change is just at the trait layer. The goal here is to decouple
changing the trait's API while preserving behavior on either side.

The `SlashCommandOutput` type now has two methods for converting two and
from a stream to use in cases where we're not yet doing streaming.

On the `SlashCommand` implementer side, the implements can call
`to_event_stream` to produce a stream of events based off the
`SlashCommandOutput`.

On the slash command consumer side we use
`SlashCommandOutput::from_event_stream` to convert a stream of events
back into a `SlashCommandOutput`.

The `/file` slash command has been updated to emit `SlashCommandEvent`s
directly in order for it to work properly.

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>
This commit is contained in:
Marshall Bowers 2024-10-23 21:26:50 -04:00 committed by GitHub
parent 510c71d41b
commit d30361537e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 516 additions and 88 deletions

View file

@ -1,13 +1,15 @@
use anyhow::{anyhow, Context as _, Result};
use assistant_slash_command::{
AfterCompletion, ArgumentCompletion, SlashCommand, SlashCommandOutput,
SlashCommandOutputSection, SlashCommandResult,
AfterCompletion, ArgumentCompletion, SlashCommand, SlashCommandContent, SlashCommandEvent,
SlashCommandOutput, SlashCommandOutputSection, SlashCommandResult,
};
use futures::channel::mpsc;
use fuzzy::PathMatch;
use gpui::{AppContext, Model, Task, View, WeakView};
use language::{BufferSnapshot, CodeLabel, HighlightId, LineEnding, LspAdapterDelegate};
use project::{PathMatchCandidateSet, Project};
use serde::{Deserialize, Serialize};
use smol::stream::StreamExt;
use std::{
fmt::Write,
ops::{Range, RangeInclusive},
@ -221,11 +223,11 @@ fn collect_files(
.map(|worktree| worktree.read(cx).snapshot())
.collect::<Vec<_>>();
let (events_tx, events_rx) = mpsc::unbounded();
cx.spawn(|mut cx| async move {
let mut output = SlashCommandOutput::default();
for snapshot in snapshots {
let worktree_id = snapshot.id();
let mut directory_stack: Vec<(Arc<Path>, String, usize)> = Vec::new();
let mut directory_stack: Vec<Arc<Path>> = Vec::new();
let mut folded_directory_names_stack = Vec::new();
let mut is_top_level_directory = true;
@ -241,17 +243,19 @@ fn collect_files(
continue;
}
while let Some((dir, _, _)) = directory_stack.last() {
while let Some(dir) = directory_stack.last() {
if entry.path.starts_with(dir) {
break;
}
let (_, entry_name, start) = directory_stack.pop().unwrap();
output.sections.push(build_entry_output_section(
start..output.text.len().saturating_sub(1),
Some(&PathBuf::from(entry_name)),
true,
None,
));
directory_stack.pop().unwrap();
events_tx
.unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
SlashCommandContent::Text {
text: "\n".into(),
run_commands_in_text: false,
},
)))?;
}
let filename = entry
@ -283,23 +287,46 @@ fn collect_files(
continue;
}
let prefix_paths = folded_directory_names_stack.drain(..).as_slice().join("/");
let entry_start = output.text.len();
if prefix_paths.is_empty() {
if is_top_level_directory {
output
.text
.push_str(&path_including_worktree_name.to_string_lossy());
let label = if is_top_level_directory {
is_top_level_directory = false;
path_including_worktree_name.to_string_lossy().to_string()
} else {
output.text.push_str(&filename);
}
directory_stack.push((entry.path.clone(), filename, entry_start));
filename
};
events_tx.unbounded_send(Ok(SlashCommandEvent::StartSection {
icon: IconName::Folder,
label: label.clone().into(),
metadata: None,
}))?;
events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
SlashCommandContent::Text {
text: label,
run_commands_in_text: false,
},
)))?;
directory_stack.push(entry.path.clone());
} else {
let entry_name = format!("{}/{}", prefix_paths, &filename);
output.text.push_str(&entry_name);
directory_stack.push((entry.path.clone(), entry_name, entry_start));
events_tx.unbounded_send(Ok(SlashCommandEvent::StartSection {
icon: IconName::Folder,
label: entry_name.clone().into(),
metadata: None,
}))?;
events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
SlashCommandContent::Text {
text: entry_name,
run_commands_in_text: false,
},
)))?;
directory_stack.push(entry.path.clone());
}
output.text.push('\n');
events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
SlashCommandContent::Text {
text: "\n".into(),
run_commands_in_text: false,
},
)))?;
} else if entry.is_file() {
let Some(open_buffer_task) = project_handle
.update(&mut cx, |project, cx| {
@ -310,6 +337,7 @@ fn collect_files(
continue;
};
if let Some(buffer) = open_buffer_task.await.log_err() {
let mut output = SlashCommandOutput::default();
let snapshot = buffer.read_with(&cx, |buffer, _| buffer.snapshot())?;
append_buffer_to_output(
&snapshot,
@ -317,32 +345,19 @@ fn collect_files(
&mut output,
)
.log_err();
let mut buffer_events = output.to_event_stream();
while let Some(event) = buffer_events.next().await {
events_tx.unbounded_send(event)?;
}
}
}
}
while let Some((dir, entry, start)) = directory_stack.pop() {
if directory_stack.is_empty() {
let mut root_path = PathBuf::new();
root_path.push(snapshot.root_name());
root_path.push(&dir);
output.sections.push(build_entry_output_section(
start..output.text.len(),
Some(&root_path),
true,
None,
));
} else {
output.sections.push(build_entry_output_section(
start..output.text.len(),
Some(&PathBuf::from(entry.as_str())),
true,
None,
));
}
while let Some(_) = directory_stack.pop() {
events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
}
}
Ok(output)
Ok(events_rx.boxed())
})
}
@ -528,8 +543,10 @@ pub fn append_buffer_to_output(
#[cfg(test)]
mod test {
use assistant_slash_command::SlashCommandOutput;
use fs::FakeFs;
use gpui::TestAppContext;
use pretty_assertions::assert_eq;
use project::Project;
use serde_json::json;
use settings::SettingsStore;
@ -577,6 +594,9 @@ mod test {
.update(|cx| collect_files(project.clone(), &["root/dir".to_string()], cx))
.await
.unwrap();
let result_1 = SlashCommandOutput::from_event_stream(result_1)
.await
.unwrap();
assert!(result_1.text.starts_with("root/dir"));
// 4 files + 2 directories
@ -586,6 +606,9 @@ mod test {
.update(|cx| collect_files(project.clone(), &["root/dir/".to_string()], cx))
.await
.unwrap();
let result_2 = SlashCommandOutput::from_event_stream(result_2)
.await
.unwrap();
assert_eq!(result_1, result_2);
@ -593,6 +616,7 @@ mod test {
.update(|cx| collect_files(project.clone(), &["root/dir*".to_string()], cx))
.await
.unwrap();
let result = SlashCommandOutput::from_event_stream(result).await.unwrap();
assert!(result.text.starts_with("root/dir"));
// 5 files + 2 directories
@ -639,6 +663,7 @@ mod test {
.update(|cx| collect_files(project.clone(), &["zed/assets/themes".to_string()], cx))
.await
.unwrap();
let result = SlashCommandOutput::from_event_stream(result).await.unwrap();
// Sanity check
assert!(result.text.starts_with("zed/assets/themes\n"));
@ -700,6 +725,7 @@ mod test {
.update(|cx| collect_files(project.clone(), &["zed/assets/themes".to_string()], cx))
.await
.unwrap();
let result = SlashCommandOutput::from_event_stream(result).await.unwrap();
assert!(result.text.starts_with("zed/assets/themes\n"));
assert_eq!(result.sections[0].label, "zed/assets/themes/LICENSE");
@ -720,6 +746,8 @@ mod test {
assert_eq!(result.sections[6].label, "summercamp");
assert_eq!(result.sections[7].label, "zed/assets/themes");
assert_eq!(result.text, "zed/assets/themes\n```zed/assets/themes/LICENSE\n1\n```\n\nsummercamp\n```zed/assets/themes/summercamp/LICENSE\n1\n```\n\nsubdir\n```zed/assets/themes/summercamp/subdir/LICENSE\n1\n```\n\nsubsubdir\n```zed/assets/themes/summercamp/subdir/subsubdir/LICENSE\n3\n```\n\n");
// Ensure that the project lasts until after the last await
drop(project);
}