assistant: Add example streaming slash command (#20034)

Marshall Bowers created 1 year ago

This PR adds a `/streaming-example` slash command for the purposes of
showcasing streaming during development.

This slash command is only available to staff and isn't intended to be
shipped to the general public.

Release Notes:

- N/A

Change summary

crates/assistant/src/assistant.rs                               |  14 
crates/assistant/src/slash_command.rs                           |   1 
crates/assistant/src/slash_command/streaming_example_command.rs | 136 +++
3 files changed, 151 insertions(+)

Detailed changes

crates/assistant/src/assistant.rs

@@ -51,6 +51,7 @@ use std::sync::Arc;
 pub(crate) use streaming_diff::*;
 use util::ResultExt;
 
+use crate::slash_command::streaming_example_command;
 use crate::slash_command_settings::SlashCommandSettings;
 
 actions!(
@@ -468,6 +469,19 @@ fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut
     })
     .detach();
 
+    cx.observe_flag::<streaming_example_command::StreamingExampleSlashCommandFeatureFlag, _>({
+        let slash_command_registry = slash_command_registry.clone();
+        move |is_enabled, _cx| {
+            if is_enabled {
+                slash_command_registry.register_command(
+                    streaming_example_command::StreamingExampleSlashCommand,
+                    false,
+                );
+            }
+        }
+    })
+    .detach();
+
     update_slash_commands_from_settings(cx);
     cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
         .detach();

crates/assistant/src/slash_command.rs

@@ -31,6 +31,7 @@ pub mod now_command;
 pub mod project_command;
 pub mod prompt_command;
 pub mod search_command;
+pub mod streaming_example_command;
 pub mod symbols_command;
 pub mod tab_command;
 pub mod terminal_command;

crates/assistant/src/slash_command/streaming_example_command.rs

@@ -0,0 +1,136 @@
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Result;
+use assistant_slash_command::{
+    ArgumentCompletion, SlashCommand, SlashCommandContent, SlashCommandEvent,
+    SlashCommandOutputSection, SlashCommandResult,
+};
+use feature_flags::FeatureFlag;
+use futures::channel::mpsc;
+use gpui::{Task, WeakView};
+use language::{BufferSnapshot, LspAdapterDelegate};
+use smol::stream::StreamExt;
+use smol::Timer;
+use ui::prelude::*;
+use workspace::Workspace;
+
+pub struct StreamingExampleSlashCommandFeatureFlag;
+
+impl FeatureFlag for StreamingExampleSlashCommandFeatureFlag {
+    const NAME: &'static str = "streaming-example-slash-command";
+}
+
+pub(crate) struct StreamingExampleSlashCommand;
+
+impl SlashCommand for StreamingExampleSlashCommand {
+    fn name(&self) -> String {
+        "streaming-example".into()
+    }
+
+    fn description(&self) -> String {
+        "An example slash command that showcases streaming.".into()
+    }
+
+    fn menu_text(&self) -> String {
+        self.description()
+    }
+
+    fn requires_argument(&self) -> bool {
+        false
+    }
+
+    fn complete_argument(
+        self: Arc<Self>,
+        _arguments: &[String],
+        _cancel: Arc<AtomicBool>,
+        _workspace: Option<WeakView<Workspace>>,
+        _cx: &mut WindowContext,
+    ) -> Task<Result<Vec<ArgumentCompletion>>> {
+        Task::ready(Ok(Vec::new()))
+    }
+
+    fn run(
+        self: Arc<Self>,
+        _arguments: &[String],
+        _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
+        _context_buffer: BufferSnapshot,
+        _workspace: WeakView<Workspace>,
+        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
+        cx: &mut WindowContext,
+    ) -> Task<SlashCommandResult> {
+        let (events_tx, events_rx) = mpsc::unbounded();
+        cx.background_executor()
+            .spawn(async move {
+                events_tx.unbounded_send(Ok(SlashCommandEvent::StartSection {
+                    icon: IconName::FileRust,
+                    label: "Section 1".into(),
+                    metadata: None,
+                }))?;
+                events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
+                    SlashCommandContent::Text {
+                        text: "Hello".into(),
+                        run_commands_in_text: false,
+                    },
+                )))?;
+                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,
+                    },
+                )))?;
+
+                Timer::after(Duration::from_secs(1)).await;
+
+                events_tx.unbounded_send(Ok(SlashCommandEvent::StartSection {
+                    icon: IconName::FileRust,
+                    label: "Section 2".into(),
+                    metadata: None,
+                }))?;
+                events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
+                    SlashCommandContent::Text {
+                        text: "World".into(),
+                        run_commands_in_text: false,
+                    },
+                )))?;
+                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,
+                    },
+                )))?;
+
+                for n in 1..=10 {
+                    Timer::after(Duration::from_secs(1)).await;
+
+                    events_tx.unbounded_send(Ok(SlashCommandEvent::StartSection {
+                        icon: IconName::StarFilled,
+                        label: format!("Section {n}").into(),
+                        metadata: None,
+                    }))?;
+                    events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
+                        SlashCommandContent::Text {
+                            text: "lorem ipsum ".repeat(n).trim().into(),
+                            run_commands_in_text: false,
+                        },
+                    )))?;
+                    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,
+                        },
+                    )))?;
+                }
+
+                anyhow::Ok(())
+            })
+            .detach_and_log_err(cx);
+
+        Task::ready(Ok(events_rx.boxed()))
+    }
+}