Skip to content

Commit b49df73

Browse files
committed
adjusts spinners to be compatible with conduit
1 parent 2d1d957 commit b49df73

File tree

5 files changed

+87
-55
lines changed

5 files changed

+87
-55
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/chat-cli-ui/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ serde.workspace = true
1919
serde_json.workspace = true
2020
chrono.workspace = true
2121
crossterm.workspace = true
22+
tokio.workspace = true
2223

2324
[target.'cfg(unix)'.dependencies]
2425
nix.workspace = true

crates/chat-cli-ui/src/conduit.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ pub enum ConduitError {
2121
pub struct ViewEnd {
2222
/// Used by the view to send input to the control
2323
// TODO: later on we will need replace this byte array with an actual event type from ACP
24-
pub sender: std::sync::mpsc::Sender<Vec<u8>>,
24+
pub sender: tokio::sync::mpsc::Sender<Vec<u8>>,
2525
/// To receive messages from control about state changes
2626
pub receiver: std::sync::mpsc::Receiver<Event>,
2727
}
@@ -61,7 +61,7 @@ pub struct DestinationStderr;
6161
#[derive(Clone, Debug)]
6262
pub struct DestinationStructuredOutput;
6363

64-
pub type InputReceiver = std::sync::mpsc::Receiver<Vec<u8>>;
64+
pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;
6565

6666
/// This compliments the [ViewEnd]. It can be thought of as the "other end" of a pipe.
6767
/// The control would own this.
@@ -212,7 +212,7 @@ pub fn get_legacy_conduits() -> (
212212
ControlEnd<DestinationStdout>,
213213
) {
214214
let (state_tx, state_rx) = std::sync::mpsc::channel::<Event>();
215-
let (byte_tx, byte_rx) = std::sync::mpsc::channel::<Vec<u8>>();
215+
let (byte_tx, byte_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(10);
216216

217217
(
218218
ViewEnd {
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
use crossterm::{cursor, execute, style, terminal};
2+
use indicatif::{ProgressBar, ProgressStyle};
3+
use tokio_util::sync::CancellationToken;
4+
5+
const SPINNER_CHARS: &str = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏";
6+
7+
pub struct Spinners {
8+
cancellation_token: CancellationToken,
9+
}
10+
11+
impl Spinners {
12+
pub fn new(message: String) -> Self {
13+
let pb = ProgressBar::new_spinner();
14+
pb.set_style(
15+
ProgressStyle::default_spinner()
16+
.tick_chars(SPINNER_CHARS)
17+
.template("{spinner:.green} {msg}")
18+
.unwrap(),
19+
);
20+
pb.set_message(message);
21+
let token = CancellationToken::new();
22+
let token_clone = token.clone();
23+
tokio::spawn(async move {
24+
loop {
25+
tokio::select! {
26+
_ = token_clone.cancelled() => {
27+
break;
28+
},
29+
_ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
30+
pb.tick();
31+
}
32+
}
33+
}
34+
35+
Ok::<(), Box<dyn std::error::Error + Send + Sync + 'static>>(())
36+
});
37+
38+
Self {
39+
cancellation_token: token,
40+
}
41+
}
42+
}
43+
44+
impl Drop for Spinners {
45+
fn drop(&mut self) {
46+
self.cancellation_token.cancel();
47+
let _ = execute!(
48+
std::io::stderr(),
49+
cursor::MoveToColumn(0),
50+
terminal::Clear(terminal::ClearType::CurrentLine),
51+
style::Print("\n"),
52+
);
53+
}
54+
}

crates/chat-cli/src/cli/chat/mod.rs

Lines changed: 28 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ pub mod cli;
33
mod consts;
44
pub mod context;
55
mod conversation;
6+
mod custom_spinner;
67
mod input_source;
78
mod message;
89
mod parse;
@@ -62,7 +63,6 @@ pub use conversation::ConversationState;
6263
use conversation::TokenWarningLevel;
6364
use crossterm::style::{
6465
Attribute,
65-
Color,
6666
Stylize,
6767
};
6868
use crossterm::{
@@ -72,6 +72,7 @@ use crossterm::{
7272
style,
7373
terminal,
7474
};
75+
use custom_spinner::Spinners;
7576
use eyre::{
7677
Report,
7778
Result,
@@ -96,10 +97,6 @@ use parser::{
9697
};
9798
use regex::Regex;
9899
use rmcp::model::PromptMessage;
99-
use spinners::{
100-
Spinner,
101-
Spinners,
102-
};
103100
use thiserror::Error;
104101
use time::OffsetDateTime;
105102
use token_counter::TokenCounter;
@@ -257,7 +254,6 @@ impl ChatArgs {
257254
}
258255
}
259256

260-
let stdout = std::io::stdout();
261257
let mut stderr = std::io::stderr();
262258

263259
let args: Vec<String> = std::env::args().collect();
@@ -420,8 +416,6 @@ impl ChatArgs {
420416

421417
ChatSession::new(
422418
os,
423-
stdout,
424-
stderr,
425419
&conversation_id,
426420
agents,
427421
input,
@@ -567,7 +561,7 @@ pub struct ChatSession {
567561
input_source: InputSource,
568562
/// Width of the terminal, required for [ParseState].
569563
terminal_width_provider: fn() -> Option<usize>,
570-
spinner: Option<Spinner>,
564+
spinner: Option<Spinners>,
571565
/// [ConversationState].
572566
conversation: ConversationState,
573567
/// Tool uses requested by the model that are actively being handled.
@@ -599,8 +593,6 @@ impl ChatSession {
599593
#[allow(clippy::too_many_arguments)]
600594
pub async fn new(
601595
os: &mut Os,
602-
stdout: std::io::Stdout,
603-
mut stderr: std::io::Stderr,
604596
conversation_id: &str,
605597
mut agents: Agents,
606598
mut input: Option<String>,
@@ -621,7 +613,7 @@ impl ChatSession {
621613
.and_then(|cwd| os.database.get_conversation_by_path(cwd).ok())
622614
.flatten();
623615

624-
let (view_end, _byte_receiver, control_end_stderr, control_end_stdout) = get_legacy_conduits();
616+
let (view_end, _byte_receiver, mut control_end_stderr, control_end_stdout) = get_legacy_conduits();
625617

626618
tokio::task::spawn_blocking(move || {
627619
let stderr = std::io::stderr();
@@ -646,7 +638,7 @@ impl ChatSession {
646638
if let Some(profile) = cs.current_profile() {
647639
if agents.switch(profile).is_err() {
648640
execute!(
649-
stderr,
641+
control_end_stderr,
650642
StyledText::error_fg(),
651643
style::Print("Error"),
652644
StyledText::reset(),
@@ -1102,10 +1094,6 @@ impl ChatSession {
11021094

11031095
impl Drop for ChatSession {
11041096
fn drop(&mut self) {
1105-
if let Some(spinner) = &mut self.spinner {
1106-
spinner.stop();
1107-
}
1108-
11091097
execute!(
11101098
self.stderr,
11111099
cursor::MoveToColumn(0),
@@ -1404,7 +1392,7 @@ impl ChatSession {
14041392

14051393
if self.interactive {
14061394
execute!(self.stderr, cursor::Hide, style::Print("\n"))?;
1407-
self.spinner = Some(Spinner::new(Spinners::Dots, "Creating summary...".to_string()));
1395+
self.spinner = Some(Spinners::new("Creating summary...".to_string()));
14081396
}
14091397

14101398
let mut response = match self
@@ -1699,10 +1687,10 @@ impl ChatSession {
16991687

17001688
if self.interactive {
17011689
execute!(self.stderr, cursor::Hide, style::Print("\n"))?;
1702-
self.spinner = Some(Spinner::new(
1703-
Spinners::Dots,
1704-
format!("Generating agent config for '{}'...", agent_name),
1705-
));
1690+
self.spinner = Some(Spinners::new(format!(
1691+
"Generating agent config for '{}'...",
1692+
agent_name
1693+
)));
17061694
}
17071695

17081696
let mut response = match self
@@ -2146,7 +2134,7 @@ impl ChatSession {
21462134
queue!(self.stderr, cursor::Hide)?;
21472135

21482136
if self.interactive {
2149-
self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_owned()));
2137+
self.spinner = Some(Spinners::new("Thinking...".to_owned()));
21502138
}
21512139

21522140
Ok(ChatState::HandleResponseStream(conv_state))
@@ -2567,7 +2555,7 @@ impl ChatSession {
25672555
execute!(self.stderr, cursor::Hide)?;
25682556
execute!(self.stderr, style::Print("\n"), StyledText::reset_attributes())?;
25692557
if self.interactive {
2570-
self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_string()));
2558+
self.spinner = Some(Spinners::new("Thinking...".to_string()));
25712559
}
25722560

25732561
self.send_chat_telemetry(os, TelemetryResult::Succeeded, None, None, None, false)
@@ -2715,7 +2703,7 @@ impl ChatSession {
27152703
);
27162704

27172705
execute!(self.stderr, cursor::Hide)?;
2718-
self.spinner = Some(Spinner::new(Spinners::Dots, "Dividing up the work...".to_string()));
2706+
self.spinner = Some(Spinners::new("Dividing up the work...".to_string()));
27192707

27202708
// For stream timeouts, we'll tell the model to try and split its response into
27212709
// smaller chunks.
@@ -2889,7 +2877,7 @@ impl ChatSession {
28892877
if tool_name_being_recvd.is_some() {
28902878
queue!(self.stderr, cursor::Hide)?;
28912879
if self.interactive {
2892-
self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_string()));
2880+
self.spinner = Some(Spinners::new("Thinking...".to_string()));
28932881
}
28942882
}
28952883

@@ -3206,7 +3194,7 @@ impl ChatSession {
32063194
}
32073195

32083196
if self.interactive {
3209-
self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_owned()));
3197+
self.spinner = Some(Spinners::new("Thinking...".to_owned()));
32103198
}
32113199

32123200
Ok(ChatState::HandleResponseStream(
@@ -3616,32 +3604,34 @@ async fn get_subscription_status(os: &mut Os) -> Result<ActualSubscriptionStatus
36163604

36173605
async fn get_subscription_status_with_spinner(
36183606
os: &mut Os,
3619-
output: &mut impl Write,
3607+
output: &mut (impl Write + Clone + Send + Sync + 'static),
36203608
) -> Result<ActualSubscriptionStatus> {
36213609
return with_spinner(output, "Checking subscription status...", || async {
36223610
get_subscription_status(os).await
36233611
})
36243612
.await;
36253613
}
36263614

3627-
pub async fn with_spinner<T, E, F, Fut>(output: &mut impl std::io::Write, spinner_text: &str, f: F) -> Result<T, E>
3615+
pub async fn with_spinner<T, E, F, Fut, S: std::io::Write + Clone + Send + Sync + 'static>(
3616+
output: &mut S,
3617+
spinner_text: &str,
3618+
f: F,
3619+
) -> Result<T, E>
36283620
where
36293621
F: FnOnce() -> Fut,
36303622
Fut: std::future::Future<Output = Result<T, E>>,
36313623
{
36323624
queue!(output, cursor::Hide,).ok();
3633-
let spinner = Some(Spinner::new(Spinners::Dots, spinner_text.to_owned()));
3625+
let spinner = Spinners::new(spinner_text.to_owned());
36343626

36353627
let result = f().await;
36363628

3637-
if let Some(mut s) = spinner {
3638-
s.stop();
3639-
let _ = queue!(
3640-
output,
3641-
terminal::Clear(terminal::ClearType::CurrentLine),
3642-
cursor::MoveToColumn(0),
3643-
);
3644-
}
3629+
drop(spinner);
3630+
let _ = queue!(
3631+
output,
3632+
terminal::Clear(terminal::ClearType::CurrentLine),
3633+
cursor::MoveToColumn(0),
3634+
);
36453635

36463636
result
36473637
}
@@ -3752,8 +3742,6 @@ mod tests {
37523742
.expect("Tools failed to load");
37533743
ChatSession::new(
37543744
&mut os,
3755-
std::io::stdout(),
3756-
std::io::stderr(),
37573745
"fake_conv_id",
37583746
agents,
37593747
None,
@@ -3882,8 +3870,6 @@ mod tests {
38823870
.expect("Tools failed to load");
38833871
ChatSession::new(
38843872
&mut os,
3885-
std::io::stdout(),
3886-
std::io::stderr(),
38873873
"fake_conv_id",
38883874
agents,
38893875
None,
@@ -3989,8 +3975,6 @@ mod tests {
39893975
.expect("Tools failed to load");
39903976
ChatSession::new(
39913977
&mut os,
3992-
std::io::stdout(),
3993-
std::io::stderr(),
39943978
"fake_conv_id",
39953979
agents,
39963980
None,
@@ -4067,8 +4051,6 @@ mod tests {
40674051
.expect("Tools failed to load");
40684052
ChatSession::new(
40694053
&mut os,
4070-
std::io::stdout(),
4071-
std::io::stderr(),
40724054
"fake_conv_id",
40734055
agents,
40744056
None,
@@ -4125,8 +4107,6 @@ mod tests {
41254107
.expect("Tools failed to load");
41264108
ChatSession::new(
41274109
&mut os,
4128-
std::io::stdout(),
4129-
std::io::stderr(),
41304110
"fake_conv_id",
41314111
agents,
41324112
None,
@@ -4230,8 +4210,6 @@ mod tests {
42304210
// Test that PreToolUse hook runs
42314211
ChatSession::new(
42324212
&mut os,
4233-
std::io::stdout(),
4234-
std::io::stderr(),
42354213
"fake_conv_id",
42364214
agents,
42374215
None, // No initial input
@@ -4366,8 +4344,6 @@ mod tests {
43664344
// Run chat session - hook should block tool execution
43674345
let result = ChatSession::new(
43684346
&mut os,
4369-
std::io::stdout(),
4370-
std::io::stderr(),
43714347
"test_conv_id",
43724348
agents,
43734349
None,

0 commit comments

Comments
 (0)