diff options
Diffstat (limited to 'src/render')
| -rw-r--r-- | src/render/mod.rs | 5 | ||||
| -rw-r--r-- | src/render/stream.rs | 12 |
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()?; } |
