1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
|
mod cmd;
mod markdown;
mod repl;
use self::cmd::cmd_render_stream;
pub use self::markdown::{MarkdownRender, RenderOptions};
use self::repl::repl_render_stream;
use crate::client::Client;
use crate::config::GlobalConfig;
use crate::utils::AbortSignal;
use anyhow::{Context, Result};
use crossbeam::channel::{unbounded, Sender};
use crossbeam::sync::WaitGroup;
use nu_ansi_term::{Color, Style};
use std::thread::spawn;
pub fn render_stream(
input: &str,
client: &dyn Client,
config: &GlobalConfig,
repl: bool,
abort: AbortSignal,
wg: WaitGroup,
) -> Result<String> {
let render_options = config.read().get_render_options()?;
let mut stream_handler = {
let (tx, rx) = unbounded();
let abort_clone = abort.clone();
let highlight = config.read().highlight;
spawn(move || {
let run = move || {
if repl {
let mut render = MarkdownRender::init(render_options)?;
repl_render_stream(&rx, &mut render, &abort)
} else {
let mut render = MarkdownRender::init(render_options)?;
cmd_render_stream(&rx, &mut render, &abort)
}
};
if let Err(err) = run() {
render_error(err, highlight);
}
drop(wg);
});
ReplyHandler::new(tx, abort_clone)
};
client.send_message_streaming(input, &mut stream_handler)?;
let buffer = stream_handler.get_buffer();
Ok(buffer.to_string())
}
pub fn render_error(err: anyhow::Error, highlight: bool) {
let err = format!("{err:?}");
if highlight {
let style = Style::new().fg(Color::Red);
println!("{}", style.paint(err.trim()));
} else {
println!("{}", err.trim());
}
}
pub struct ReplyHandler {
sender: Sender<ReplyEvent>,
buffer: String,
abort: AbortSignal,
}
impl ReplyHandler {
pub fn new(sender: Sender<ReplyEvent>, abort: AbortSignal) -> Self {
Self {
sender,
abort,
buffer: String::new(),
}
}
pub fn text(&mut self, text: &str) -> Result<()> {
if self.buffer.is_empty() && text == "\n\n" {
return Ok(());
}
self.buffer.push_str(text);
let ret = self
.sender
.send(ReplyEvent::Text(text.to_string()))
.with_context(|| "Failed to send ReplyEvent:Text");
self.safe_ret(ret)?;
Ok(())
}
pub fn done(&mut self) -> Result<()> {
let ret = self
.sender
.send(ReplyEvent::Done)
.with_context(|| "Failed to send ReplyEvent::Done");
self.safe_ret(ret)?;
Ok(())
}
pub fn get_buffer(&self) -> &str {
&self.buffer
}
pub fn get_abort(&self) -> AbortSignal {
self.abort.clone()
}
fn safe_ret(&self, ret: Result<()>) -> Result<()> {
if ret.is_err() && self.abort.aborted() {
return Ok(());
}
ret
}
}
pub enum ReplyEvent {
Text(String),
Done,
}
|
