assistant: Update SlashCommand trait with streaming return type (#1…

…9652) 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 <[email protected]>
zed-industries · Oct 24, 2024 · d303615 · d303615
1 parent 510c71d
commit d303615
Show file tree

Hide file tree

Showing 23 changed files with 516 additions and 88 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/assistant/src/context.rs b/crates/assistant/src/context.rs
@@ -7,7 +7,7 @@ use crate::{
 };
 use anyhow::{anyhow, Context as _, Result};
 use assistant_slash_command::{
-    SlashCommandOutputSection, SlashCommandRegistry, SlashCommandResult,
+    SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry, SlashCommandResult,
 };
 use assistant_tool::ToolRegistry;
 use client::{self, proto, telemetry::Telemetry};
@@ -1688,19 +1688,13 @@ impl Context {
             let command_range = command_range.clone();
             async move {
                 let output = output.await;
+                let output = match output {
+                    Ok(output) => SlashCommandOutput::from_event_stream(output).await,
+                    Err(err) => Err(err),
+                };
                 this.update(&mut cx, |this, cx| match output {
                     Ok(mut output) => {
-                        // Ensure section ranges are valid.
-                        for section in &mut output.sections {
-                            section.range.start = section.range.start.min(output.text.len());
-                            section.range.end = section.range.end.min(output.text.len());
-                            while !output.text.is_char_boundary(section.range.start) {
-                                section.range.start -= 1;
-                            }
-                            while !output.text.is_char_boundary(section.range.end) {
-                                section.range.end += 1;
-                            }
-                        }
+                        output.ensure_valid_section_ranges();
 
                         // Ensure there is a newline after the last section.
                         if ensure_trailing_newline {

diff --git a/crates/assistant/src/context/context_tests.rs b/crates/assistant/src/context/context_tests.rs
@@ -1097,7 +1097,8 @@ async fn test_random_context_collaboration(cx: &mut TestAppContext, mut rng: Std
                             text: output_text,
                             sections,
                             run_commands_in_text: false,
-                        })),
+                        }
+                        .to_event_stream())),
                         true,
                         false,
                         cx,
@@ -1421,6 +1422,7 @@ impl SlashCommand for FakeSlashCommand {
             text: format!("Executed fake command: {}", self.0),
             sections: vec![],
             run_commands_in_text: false,
-        }))
+        }
+        .to_event_stream()))
     }
 }
diff --git a/crates/assistant/src/slash_command/auto_command.rs b/crates/assistant/src/slash_command/auto_command.rs
@@ -147,7 +147,8 @@ impl SlashCommand for AutoCommand {
                 text: prompt,
                 sections: Vec::new(),
                 run_commands_in_text: true,
-            })
+            }
+            .to_event_stream())
         })
     }
 }

diff --git a/crates/assistant/src/slash_command/cargo_workspace_command.rs b/crates/assistant/src/slash_command/cargo_workspace_command.rs
@@ -147,7 +147,8 @@ impl SlashCommand for CargoWorkspaceSlashCommand {
                         metadata: None,
                     }],
                     run_commands_in_text: false,
-                })
+                }
+                .to_event_stream())
             })
         });
         output.unwrap_or_else(|error| Task::ready(Err(error)))

diff --git a/crates/assistant/src/slash_command/context_server_command.rs b/crates/assistant/src/slash_command/context_server_command.rs
@@ -185,7 +185,8 @@ impl SlashCommand for ContextServerSlashCommand {
                     }],
                     text: prompt,
                     run_commands_in_text: false,
-                })
+                }
+                .to_event_stream())
             })
         } else {
             Task::ready(Err(anyhow!("Context server not found")))

diff --git a/crates/assistant/src/slash_command/default_command.rs b/crates/assistant/src/slash_command/default_command.rs
@@ -78,7 +78,8 @@ impl SlashCommand for DefaultSlashCommand {
                 }],
                 text,
                 run_commands_in_text: true,
-            })
+            }
+            .to_event_stream())
         })
     }
 }
diff --git a/crates/assistant/src/slash_command/delta_command.rs b/crates/assistant/src/slash_command/delta_command.rs
@@ -86,25 +86,28 @@ impl SlashCommand for DeltaSlashCommand {
                 .zip(file_command_new_outputs)
             {
                 if let Ok(new_output) = new_output {
-                    if let Some(file_command_range) = new_output.sections.first() {
-                        let new_text = &new_output.text[file_command_range.range.clone()];
-                        if old_text.chars().ne(new_text.chars()) {
-                            output.sections.extend(new_output.sections.into_iter().map(
-                                |section| SlashCommandOutputSection {
-                                    range: output.text.len() + section.range.start
-                                        ..output.text.len() + section.range.end,
-                                    icon: section.icon,
-                                    label: section.label,
-                                    metadata: section.metadata,
-                                },
-                            ));
-                            output.text.push_str(&new_output.text);
+                    if let Ok(new_output) = SlashCommandOutput::from_event_stream(new_output).await
+                    {
+                        if let Some(file_command_range) = new_output.sections.first() {
+                            let new_text = &new_output.text[file_command_range.range.clone()];
+                            if old_text.chars().ne(new_text.chars()) {
+                                output.sections.extend(new_output.sections.into_iter().map(
+                                    |section| SlashCommandOutputSection {
+                                        range: output.text.len() + section.range.start
+                                            ..output.text.len() + section.range.end,
+                                        icon: section.icon,
+                                        label: section.label,
+                                        metadata: section.metadata,
+                                    },
+                                ));
+                                output.text.push_str(&new_output.text);
+                            }
                         }
                     }
                 }
             }
 
-            Ok(output)
+            Ok(output.to_event_stream())
         })
     }
 }
diff --git a/crates/assistant/src/slash_command/diagnostics_command.rs b/crates/assistant/src/slash_command/diagnostics_command.rs
@@ -180,7 +180,11 @@ impl SlashCommand for DiagnosticsSlashCommand {
 
         let task = collect_diagnostics(workspace.read(cx).project().clone(), options, cx);
 
-        cx.spawn(move |_| async move { task.await?.ok_or_else(|| anyhow!("No diagnostics found")) })
+        cx.spawn(move |_| async move {
+            task.await?
+                .map(|output| output.to_event_stream())
+                .ok_or_else(|| anyhow!("No diagnostics found"))
+        })
     }
 }
 

diff --git a/crates/assistant/src/slash_command/docs_command.rs b/crates/assistant/src/slash_command/docs_command.rs
@@ -356,7 +356,8 @@ impl SlashCommand for DocsSlashCommand {
                     })
                     .collect(),
                 run_commands_in_text: false,
-            })
+            }
+            .to_event_stream())
         })
     }
 }

diff --git a/crates/assistant/src/slash_command/fetch_command.rs b/crates/assistant/src/slash_command/fetch_command.rs
@@ -167,7 +167,8 @@ impl SlashCommand for FetchSlashCommand {
                     metadata: None,
                 }],
                 run_commands_in_text: false,
-            })
+            }
+            .to_event_stream())
         })
     }
 }