Fix regression in producing sections when converting `SlashCommandOutput` to event stream (#20404)

Antonio Scandurra created 1 year ago

Closes #20243 

Release Notes:

- N/A

Change summary

crates/assistant/src/context.rs                                 |  4 
crates/assistant/src/context/context_tests.rs                   |  4 
crates/assistant/src/slash_command/file_command.rs              |  5 
crates/assistant/src/slash_command/selection_command.rs         |  2 
crates/assistant/src/slash_command/streaming_example_command.rs |  7 
crates/assistant_slash_command/src/assistant_slash_command.rs   | 79 +-
6 files changed, 41 insertions(+), 60 deletions(-)

Detailed changes

crates/assistant/src/context.rs

@@ -1983,7 +1983,7 @@ impl Context {
                                     run_commands_in_ranges.push(start..end);
                                 }
                             }
-                            SlashCommandEvent::EndSection { metadata } => {
+                            SlashCommandEvent::EndSection => {
                                 if let Some(pending_section) = pending_section_stack.pop() {
                                     let offset_range = (pending_section.start..insert_position)
                                         .to_offset(this.buffer.read(cx));
@@ -1997,7 +1997,7 @@ impl Context {
                                                 range: range.clone(),
                                                 icon: pending_section.icon,
                                                 label: pending_section.label,
-                                                metadata: metadata.or(pending_section.metadata),
+                                                metadata: pending_section.metadata,
                                             },
                                             cx,
                                         );

crates/assistant/src/context/context_tests.rs

@@ -583,7 +583,7 @@ async fn test_slash_commands(cx: &mut TestAppContext) {
     );
 
     command_output_tx
-        .unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))
+        .unbounded_send(Ok(SlashCommandEvent::EndSection))
         .unwrap();
     cx.run_until_parked();
     assert_text_and_context_ranges(
@@ -1310,7 +1310,7 @@ async fn test_random_context_collaboration(cx: &mut TestAppContext, mut rng: Std
                             text: output_text[section_start..section_end].to_string(),
                             run_commands_in_text: false,
                         })));
-                        events.push(Ok(SlashCommandEvent::EndSection { metadata: None }));
+                        events.push(Ok(SlashCommandEvent::EndSection));
                         section_start = section_end;
                     }
 

crates/assistant/src/slash_command/file_command.rs

@@ -256,8 +256,7 @@ fn collect_files(
                         break;
                     }
                     directory_stack.pop().unwrap();
-                    events_tx
-                        .unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
+                    events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection))?;
                     events_tx.unbounded_send(Ok(SlashCommandEvent::Content(
                         SlashCommandContent::Text {
                             text: "\n".into(),
@@ -362,7 +361,7 @@ fn collect_files(
             }
 
             while let Some(_) = directory_stack.pop() {
-                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
+                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection))?;
             }
         }
 

crates/assistant/src/slash_command/selection_command.rs

@@ -84,7 +84,7 @@ impl SlashCommand for SelectionCommand {
                 text,
                 run_commands_in_text: false,
             })));
-            events.push(Ok(SlashCommandEvent::EndSection { metadata: None }));
+            events.push(Ok(SlashCommandEvent::EndSection));
             events.push(Ok(SlashCommandEvent::Content(SlashCommandContent::Text {
                 text: "\n".to_string(),
                 run_commands_in_text: false,

crates/assistant/src/slash_command/streaming_example_command.rs

@@ -74,7 +74,7 @@ impl SlashCommand for StreamingExampleSlashCommand {
                         run_commands_in_text: false,
                     },
                 )))?;
-                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
+                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection))?;
 
                 Timer::after(Duration::from_secs(1)).await;
 
@@ -89,7 +89,7 @@ impl SlashCommand for StreamingExampleSlashCommand {
                         run_commands_in_text: false,
                     },
                 )))?;
-                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
+                events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection))?;
 
                 for n in 1..=10 {
                     Timer::after(Duration::from_secs(1)).await;
@@ -105,8 +105,7 @@ impl SlashCommand for StreamingExampleSlashCommand {
                             run_commands_in_text: false,
                         },
                     )))?;
-                    events_tx
-                        .unbounded_send(Ok(SlashCommandEvent::EndSection { metadata: None }))?;
+                    events_tx.unbounded_send(Ok(SlashCommandEvent::EndSection))?;
                 }
 
                 anyhow::Ok(())

crates/assistant_slash_command/src/assistant_slash_command.rs

@@ -133,9 +133,7 @@ pub enum SlashCommandEvent {
         metadata: Option<serde_json::Value>,
     },
     Content(SlashCommandContent),
-    EndSection {
-        metadata: Option<serde_json::Value>,
-    },
+    EndSection,
 }
 
 #[derive(Debug, Default, PartialEq, Clone)]
@@ -164,43 +162,37 @@ impl SlashCommandOutput {
         self.ensure_valid_section_ranges();
 
         let mut events = Vec::new();
-        let mut last_section_end = 0;
 
+        let mut section_endpoints = Vec::new();
         for section in self.sections {
-            if last_section_end < section.range.start {
+            section_endpoints.push((
+                section.range.start,
+                SlashCommandEvent::StartSection {
+                    icon: section.icon,
+                    label: section.label,
+                    metadata: section.metadata,
+                },
+            ));
+            section_endpoints.push((section.range.end, SlashCommandEvent::EndSection));
+        }
+        section_endpoints.sort_by_key(|(offset, _)| *offset);
+
+        let mut content_offset = 0;
+        for (endpoint_offset, endpoint) in section_endpoints {
+            if content_offset < endpoint_offset {
                 events.push(Ok(SlashCommandEvent::Content(SlashCommandContent::Text {
-                    text: self
-                        .text
-                        .get(last_section_end..section.range.start)
-                        .unwrap_or_default()
-                        .to_string(),
+                    text: self.text[content_offset..endpoint_offset].to_string(),
                     run_commands_in_text: self.run_commands_in_text,
                 })));
+                content_offset = endpoint_offset;
             }
 
-            events.push(Ok(SlashCommandEvent::StartSection {
-                icon: section.icon,
-                label: section.label,
-                metadata: section.metadata.clone(),
-            }));
-            events.push(Ok(SlashCommandEvent::Content(SlashCommandContent::Text {
-                text: self
-                    .text
-                    .get(section.range.start..section.range.end)
-                    .unwrap_or_default()
-                    .to_string(),
-                run_commands_in_text: self.run_commands_in_text,
-            })));
-            events.push(Ok(SlashCommandEvent::EndSection {
-                metadata: section.metadata,
-            }));
-
-            last_section_end = section.range.end;
+            events.push(Ok(endpoint));
         }
 
-        if last_section_end < self.text.len() {
+        if content_offset < self.text.len() {
             events.push(Ok(SlashCommandEvent::Content(SlashCommandContent::Text {
-                text: self.text[last_section_end..].to_string(),
+                text: self.text[content_offset..].to_string(),
                 run_commands_in_text: self.run_commands_in_text,
             })));
         }
@@ -240,9 +232,8 @@ impl SlashCommandOutput {
                         section.range.end = output.text.len();
                     }
                 }
-                SlashCommandEvent::EndSection { metadata } => {
-                    if let Some(mut section) = section_stack.pop() {
-                        section.metadata = metadata;
+                SlashCommandEvent::EndSection => {
+                    if let Some(section) = section_stack.pop() {
                         output.sections.push(section);
                     }
                 }
@@ -314,7 +305,7 @@ mod tests {
                         text: "Hello, world!".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection { metadata: None }
+                    SlashCommandEvent::EndSection
                 ]
             );
 
@@ -366,7 +357,7 @@ mod tests {
                         text: "Apple\n".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection { metadata: None },
+                    SlashCommandEvent::EndSection,
                     SlashCommandEvent::Content(SlashCommandContent::Text {
                         text: "Cucumber\n".into(),
                         run_commands_in_text: false
@@ -380,7 +371,7 @@ mod tests {
                         text: "Banana\n".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection { metadata: None }
+                    SlashCommandEvent::EndSection
                 ]
             );
 
@@ -444,9 +435,7 @@ mod tests {
                         text: "Line 1".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection {
-                        metadata: Some(json!({ "a": true }))
-                    },
+                    SlashCommandEvent::EndSection,
                     SlashCommandEvent::Content(SlashCommandContent::Text {
                         text: "\n".into(),
                         run_commands_in_text: false
@@ -460,9 +449,7 @@ mod tests {
                         text: "Line 2".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection {
-                        metadata: Some(json!({ "b": true }))
-                    },
+                    SlashCommandEvent::EndSection,
                     SlashCommandEvent::Content(SlashCommandContent::Text {
                         text: "\n".into(),
                         run_commands_in_text: false
@@ -476,9 +463,7 @@ mod tests {
                         text: "Line 3".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection {
-                        metadata: Some(json!({ "c": true }))
-                    },
+                    SlashCommandEvent::EndSection,
                     SlashCommandEvent::Content(SlashCommandContent::Text {
                         text: "\n".into(),
                         run_commands_in_text: false
@@ -492,9 +477,7 @@ mod tests {
                         text: "Line 4".into(),
                         run_commands_in_text: false
                     }),
-                    SlashCommandEvent::EndSection {
-                        metadata: Some(json!({ "d": true }))
-                    },
+                    SlashCommandEvent::EndSection,
                     SlashCommandEvent::Content(SlashCommandContent::Text {
                         text: "\n".into(),
                         run_commands_in_text: false