Skip to main content

codei_tui/
repl.rs

1use std::io::{self, Write};
2use std::sync::{Arc, RwLock};
3
4use anyhow::Result;
5use codei_agent::{AgentEvent, AgentLoop};
6use codei_commands::{parse_input, Input};
7use codei_i18n::t_fmt;
8use codei_llm::Usage;
9use codei_tools::{handler_for_policy, ApprovalPolicy, ToolContext};
10use tokio::sync::mpsc;
11
12use crate::launch::InteractiveLaunch;
13use crate::slash::{handle_slash, SlashContext};
14
15pub struct ReplOptions {
16    pub auto_approve: bool,
17}
18
19pub async fn run_repl(launch: InteractiveLaunch, opts: ReplOptions) -> Result<()> {
20    let InteractiveLaunch {
21        config,
22        provider,
23        provider_name,
24        model,
25        mut session,
26        store,
27        mcp,
28    } = launch;
29    let policy = if opts.auto_approve {
30        ApprovalPolicy::Never
31    } else {
32        ApprovalPolicy::OnDestructive
33    };
34    let (tx, mut rx) = mpsc::unbounded_channel();
35    let tool_ctx = ToolContext {
36        cwd: config.cwd.clone(),
37        config: Arc::clone(&config),
38        approval: Arc::from(handler_for_policy(policy)),
39    };
40    let provider_name = Arc::new(RwLock::new(provider_name));
41    let agent = AgentLoop::new(
42        Arc::clone(&config),
43        Arc::clone(&model),
44        provider,
45        provider_name.read().expect("provider lock").clone(),
46        tool_ctx,
47        mcp,
48        Some(tx),
49    );
50
51    let stdin = io::stdin();
52    let mut stdout = io::stdout();
53    let mut token_usage = Usage::default();
54    let mut last_turn_usage = None;
55
56    loop {
57        while let Ok(event) = rx.try_recv() {
58            render_event(&mut stdout, &event, &mut token_usage, &mut last_turn_usage)?;
59        }
60
61        write!(stdout, "\n> ")?;
62        stdout.flush()?;
63        let mut line = String::new();
64        if stdin.read_line(&mut line)? == 0 {
65            break;
66        }
67
68        match parse_input(&line) {
69            Input::SlashCommand(cmd) => {
70                let mut ctx = SlashContext {
71                    session: &mut session,
72                    store: &store,
73                    model: &model,
74                    provider_name: &provider_name,
75                    agent: &agent,
76                    token_usage: &mut token_usage,
77                    last_turn_usage: &mut last_turn_usage,
78                };
79                match handle_slash(cmd, &mut ctx).await? {
80                    crate::slash::SlashAction::Exit => break,
81                    crate::slash::SlashAction::Message(text) => writeln!(stdout, "{text}")?,
82                    crate::slash::SlashAction::Continue => {}
83                }
84            }
85            Input::UserMessage(msg) if msg.is_empty() => {}
86            Input::UserMessage(msg) => {
87                if let Err(err) = agent.run_turn(&mut session, &msg, store.as_ref()).await {
88                    writeln!(
89                        stdout,
90                        "{}",
91                        t_fmt("tui_error_prefix", &[("message", &err.to_string())])
92                    )?;
93                }
94                while let Ok(event) = rx.try_recv() {
95                    render_event(&mut stdout, &event, &mut token_usage, &mut last_turn_usage)?;
96                }
97            }
98        }
99    }
100
101    Ok(())
102}
103
104fn render_event(
105    stdout: &mut impl Write,
106    event: &AgentEvent,
107    token_usage: &mut Usage,
108    last_turn_usage: &mut Option<Usage>,
109) -> io::Result<()> {
110    match event {
111        AgentEvent::AssistantDelta { text } => write!(stdout, "{text}"),
112        AgentEvent::ToolStarted { name, args } => writeln!(stdout, "\n[tool:{name}] {args}"),
113        AgentEvent::ToolFinished { name, result } => {
114            writeln!(stdout, "[tool:{name}]\n{}", result.content)
115        }
116        AgentEvent::TurnComplete { usage } => {
117            if let Some(u) = *usage {
118                token_usage.add_assign(u);
119                *last_turn_usage = Some(u);
120                writeln!(
121                    stdout,
122                    "\n{}",
123                    t_fmt(
124                        "slash_tokens_turn",
125                        &[
126                            ("input", &u.input_tokens.to_string()),
127                            ("output", &u.output_tokens.to_string()),
128                        ],
129                    )
130                )
131            } else {
132                writeln!(stdout)
133            }
134        }
135        AgentEvent::Error { message } => writeln!(
136            stdout,
137            "{}",
138            t_fmt("tui_error_prefix", &[("message", message)])
139        ),
140    }
141}