1 files changed, 27 insertions, 29 deletions

diff --git a/src/main.rs b/src/main.rs
index 368779e..dd9f9cf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,20 +10,17 @@ extern crate log;
 mod utils;
 
 use crate::cli::Cli;
-use crate::client::Client;
 use crate::config::{Config, GlobalConfig};
 
 use anyhow::Result;
 use clap::Parser;
 use client::{init_client, list_models};
-use crossbeam::sync::WaitGroup;
 use is_terminal::IsTerminal;
 use parking_lot::RwLock;
-use render::{render_stream, MarkdownRender};
+use render::{render_error, render_stream, MarkdownRender};
 use repl::Repl;
-use std::io::{stdin, Read};
+use std::io::{stderr, stdin, stdout, Read};
 use std::sync::Arc;
-use std::{io::stdout, process::exit};
 use utils::{cl100k_base_singleton, create_abort_signal};
 
 fn main() -> Result<()> {
@@ -36,18 +33,18 @@ fn main() -> Result<()> {
             .roles
             .iter()
             .for_each(|v| println!("{}", v.name));
-        exit(0);
+        return Ok(());
     }
     if cli.list_models {
         for model in list_models(&config.read()) {
             println!("{}", model.id());
         }
-        exit(0);
+        return Ok(());
     }
     if cli.list_sessions {
         let sessions = config.read().list_sessions().join("\n");
         println!("{sessions}");
-        exit(0);
+        return Ok(());
     }
     if let Some(wrap) = &cli.wrap {
         config.write().set_wrap(wrap)?;
@@ -75,15 +72,22 @@ fn main() -> Result<()> {
     if cli.info {
         let info = config.read().info()?;
         println!("{}", info);
-        exit(0);
+        return Ok(());
     }
     config.write().onstart()?;
     let no_stream = cli.no_stream;
-    let client = init_client(&config)?;
+    if let Err(err) = start(&config, text, no_stream) {
+        let highlight = stderr().is_terminal() && config.read().highlight;
+        render_error(err, highlight)
+    }
+    Ok(())
+}
+
+fn start(config: &GlobalConfig, text: Option<String>, no_stream: bool) -> Result<()> {
     if stdin().is_terminal() {
         match text {
-            Some(text) => start_directive(client.as_ref(), &config, &text, no_stream),
-            None => start_interactive(&config),
+            Some(text) => start_directive(config, &text, no_stream),
+            None => start_interactive(config),
         }
     } else {
         let mut input = String::new();
@@ -91,40 +95,34 @@ fn main() -> Result<()> {
         if let Some(text) = text {
             input = format!("{text}\n{input}");
         }
-        start_directive(client.as_ref(), &config, &input, no_stream)
+        start_directive(config, &input, no_stream)
     }
 }
 
-fn start_directive(
-    client: &dyn Client,
-    config: &GlobalConfig,
-    input: &str,
-    no_stream: bool,
-) -> Result<()> {
+fn start_directive(config: &GlobalConfig, input: &str, no_stream: bool) -> Result<()> {
     if let Some(session) = &config.read().session {
         session.guard_save()?;
     }
-    if !stdout().is_terminal() {
-        config.write().highlight = false;
-    }
+    let client = init_client(config)?;
     config.read().maybe_print_send_tokens(input);
     let output = if no_stream {
-        let render_options = config.read().get_render_options()?;
         let output = client.send_message(input)?;
-        let mut markdown_render = MarkdownRender::init(render_options)?;
-        println!("{}", markdown_render.render(&output).trim());
+        if stdout().is_terminal() {
+            let render_options = config.read().get_render_options()?;
+            let mut markdown_render = MarkdownRender::init(render_options)?;
+            println!("{}", markdown_render.render(&output).trim());
+        } else {
+            println!("{}", output);
+        }
         output
     } else {
-        let wg = WaitGroup::new();
         let abort = create_abort_signal();
         let abort_clone = abort.clone();
         ctrlc::set_handler(move || {
             abort_clone.set_ctrlc();
         })
         .expect("Failed to setting Ctrl-C handler");
-        let output = render_stream(input, client, config, false, abort, wg.clone())?;
-        wg.wait();
-        output
+        render_stream(input, client.as_ref(), config, abort)?
     };
     config.write().save_message(input, &output)
 }