Fixed subquery tool double addingmain

2 files changed, 13 insertions, 4 deletions

diff --git a/src/render/mod.rs b/src/render/mod.rs
index 0aa3184..37e8014 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -14,13 +14,14 @@ pub async fn render_stream(
     rx: UnboundedReceiver<SseEvent>,
     config: &GlobalConfig,
     abort_signal: AbortSignal,
+    prefix: Option<String>,
 ) -> Result<()> {
     let ret = if *IS_STDOUT_TERMINAL && config.read().highlight {
         let render_options = config.read().render_options()?;
         let mut render = MarkdownRender::init(render_options)?;
-        markdown_stream(rx, &mut render, &abort_signal).await
+        markdown_stream(rx, &mut render, &abort_signal, prefix).await
     } else {
-        raw_stream(rx, &abort_signal).await
+        raw_stream(rx, &abort_signal, prefix).await
     };
     ret.map_err(|err| err.context("Failed to reader stream"))
 }
diff --git a/src/render/stream.rs b/src/render/stream.rs
index af69176..e5db8b2 100644
--- a/src/render/stream.rs
+++ b/src/render/stream.rs
@@ -1,6 +1,6 @@
 use super::{MarkdownRender, SseEvent};
 
-use crate::utils::{poll_abort_signal, spawn_spinner, AbortSignal};
+use crate::utils::{dimmed_text, poll_abort_signal, spawn_spinner, AbortSignal};
 
 use anyhow::Result;
 use crossterm::{
@@ -18,7 +18,11 @@ pub async fn markdown_stream(
     rx: UnboundedReceiver<SseEvent>,
     render: &mut MarkdownRender,
     abort_signal: &AbortSignal,
+    prefix: Option<String>,
 ) -> Result<()> {
+    if let Some(prefix) = &prefix {
+        println!("{}", dimmed_text(prefix));
+    }
     enable_raw_mode()?;
     let mut stdout = io::stdout();
 
@@ -35,12 +39,13 @@ pub async fn markdown_stream(
 pub async fn raw_stream(
     mut rx: UnboundedReceiver<SseEvent>,
     abort_signal: &AbortSignal,
+    prefix: Option<String>,
 ) -> Result<()> {
     let mut spinner = Some(spawn_spinner(&match crate::utils::current_query_scope() {
         Some(scope) => format!("{scope} Generating"),
         None => "Generating".to_string(),
     }));
-
+    let mut prefix = prefix;
     loop {
         if abort_signal.aborted() {
             break;
@@ -52,6 +57,9 @@ pub async fn raw_stream(
 
             match evt {
                 SseEvent::Text(text) => {
+                    if let Some(prefix) = prefix.take() {
+                        println!("{}", dimmed_text(&prefix));
+                    }
                     print!("{text}");
                     stdout().flush()?;
                 }