input.rs 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330
  1. use std::borrow::Cow;
  2. use std::cell::RefCell;
  3. use std::collections::BTreeSet;
  4. use std::io::{self, IsTerminal, Write};
  5. use rustyline::completion::{Completer, Pair};
  6. use rustyline::error::ReadlineError;
  7. use rustyline::highlight::{CmdKind, Highlighter};
  8. use rustyline::hint::Hinter;
  9. use rustyline::history::DefaultHistory;
  10. use rustyline::validate::Validator;
  11. use rustyline::{
  12. Cmd, CompletionType, Config, Context, EditMode, Editor, Helper, KeyCode, KeyEvent, Modifiers,
  13. };
  14. #[derive(Debug, Clone, PartialEq, Eq)]
  15. pub enum ReadOutcome {
  16. Submit(String),
  17. Cancel,
  18. Exit,
  19. }
  20. struct SlashCommandHelper {
  21. completions: Vec<String>,
  22. current_line: RefCell<String>,
  23. }
  24. impl SlashCommandHelper {
  25. fn new(completions: Vec<String>) -> Self {
  26. Self {
  27. completions: normalize_completions(completions),
  28. current_line: RefCell::new(String::new()),
  29. }
  30. }
  31. fn reset_current_line(&self) {
  32. self.current_line.borrow_mut().clear();
  33. }
  34. fn current_line(&self) -> String {
  35. self.current_line.borrow().clone()
  36. }
  37. fn set_current_line(&self, line: &str) {
  38. let mut current = self.current_line.borrow_mut();
  39. current.clear();
  40. current.push_str(line);
  41. }
  42. fn set_completions(&mut self, completions: Vec<String>) {
  43. self.completions = normalize_completions(completions);
  44. }
  45. }
  46. impl Completer for SlashCommandHelper {
  47. type Candidate = Pair;
  48. fn complete(
  49. &self,
  50. line: &str,
  51. pos: usize,
  52. _ctx: &Context<'_>,
  53. ) -> rustyline::Result<(usize, Vec<Self::Candidate>)> {
  54. let Some(prefix) = slash_command_prefix(line, pos) else {
  55. return Ok((0, Vec::new()));
  56. };
  57. let matches = self
  58. .completions
  59. .iter()
  60. .filter(|candidate| candidate.starts_with(prefix))
  61. .map(|candidate| Pair {
  62. display: candidate.clone(),
  63. replacement: candidate.clone(),
  64. })
  65. .collect();
  66. Ok((0, matches))
  67. }
  68. }
  69. impl Hinter for SlashCommandHelper {
  70. type Hint = String;
  71. }
  72. impl Highlighter for SlashCommandHelper {
  73. fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
  74. self.set_current_line(line);
  75. Cow::Borrowed(line)
  76. }
  77. fn highlight_char(&self, line: &str, _pos: usize, _kind: CmdKind) -> bool {
  78. self.set_current_line(line);
  79. false
  80. }
  81. }
  82. impl Validator for SlashCommandHelper {}
  83. impl Helper for SlashCommandHelper {}
  84. pub struct LineEditor {
  85. prompt: String,
  86. editor: Editor<SlashCommandHelper, DefaultHistory>,
  87. }
  88. impl LineEditor {
  89. #[must_use]
  90. pub fn new(prompt: impl Into<String>, completions: Vec<String>) -> Self {
  91. let config = Config::builder()
  92. .completion_type(CompletionType::List)
  93. .edit_mode(EditMode::Emacs)
  94. .build();
  95. let mut editor = Editor::<SlashCommandHelper, DefaultHistory>::with_config(config)
  96. .expect("rustyline editor should initialize");
  97. editor.set_helper(Some(SlashCommandHelper::new(completions)));
  98. editor.bind_sequence(KeyEvent(KeyCode::Char('J'), Modifiers::CTRL), Cmd::Newline);
  99. editor.bind_sequence(KeyEvent(KeyCode::Enter, Modifiers::SHIFT), Cmd::Newline);
  100. Self {
  101. prompt: prompt.into(),
  102. editor,
  103. }
  104. }
  105. pub fn push_history(&mut self, entry: impl Into<String>) {
  106. let entry = entry.into();
  107. if entry.trim().is_empty() {
  108. return;
  109. }
  110. let _ = self.editor.add_history_entry(entry);
  111. }
  112. pub fn set_completions(&mut self, completions: Vec<String>) {
  113. if let Some(helper) = self.editor.helper_mut() {
  114. helper.set_completions(completions);
  115. }
  116. }
  117. pub fn read_line(&mut self) -> io::Result<ReadOutcome> {
  118. if !io::stdin().is_terminal() || !io::stdout().is_terminal() {
  119. return self.read_line_fallback();
  120. }
  121. if let Some(helper) = self.editor.helper_mut() {
  122. helper.reset_current_line();
  123. }
  124. match self.editor.readline(&self.prompt) {
  125. Ok(line) => Ok(ReadOutcome::Submit(line)),
  126. Err(ReadlineError::Interrupted) => {
  127. let has_input = !self.current_line().is_empty();
  128. self.finish_interrupted_read()?;
  129. if has_input {
  130. Ok(ReadOutcome::Cancel)
  131. } else {
  132. Ok(ReadOutcome::Exit)
  133. }
  134. }
  135. Err(ReadlineError::Eof) => {
  136. self.finish_interrupted_read()?;
  137. Ok(ReadOutcome::Exit)
  138. }
  139. Err(error) => Err(io::Error::other(error)),
  140. }
  141. }
  142. fn current_line(&self) -> String {
  143. self.editor
  144. .helper()
  145. .map_or_else(String::new, SlashCommandHelper::current_line)
  146. }
  147. fn finish_interrupted_read(&mut self) -> io::Result<()> {
  148. if let Some(helper) = self.editor.helper_mut() {
  149. helper.reset_current_line();
  150. }
  151. let mut stdout = io::stdout();
  152. writeln!(stdout)
  153. }
  154. fn read_line_fallback(&self) -> io::Result<ReadOutcome> {
  155. let mut stdout = io::stdout();
  156. write!(stdout, "{}", self.prompt)?;
  157. stdout.flush()?;
  158. let mut buffer = String::new();
  159. let bytes_read = io::stdin().read_line(&mut buffer)?;
  160. if bytes_read == 0 {
  161. return Ok(ReadOutcome::Exit);
  162. }
  163. while matches!(buffer.chars().last(), Some('\n' | '\r')) {
  164. buffer.pop();
  165. }
  166. Ok(ReadOutcome::Submit(buffer))
  167. }
  168. }
  169. fn slash_command_prefix(line: &str, pos: usize) -> Option<&str> {
  170. if pos != line.len() {
  171. return None;
  172. }
  173. let prefix = &line[..pos];
  174. if !prefix.starts_with('/') {
  175. return None;
  176. }
  177. Some(prefix)
  178. }
  179. fn normalize_completions(completions: Vec<String>) -> Vec<String> {
  180. let mut seen = BTreeSet::new();
  181. completions
  182. .into_iter()
  183. .filter(|candidate| candidate.starts_with('/'))
  184. .filter(|candidate| seen.insert(candidate.clone()))
  185. .collect()
  186. }
  187. #[cfg(test)]
  188. mod tests {
  189. use super::{slash_command_prefix, LineEditor, SlashCommandHelper};
  190. use rustyline::completion::Completer;
  191. use rustyline::highlight::Highlighter;
  192. use rustyline::history::{DefaultHistory, History};
  193. use rustyline::Context;
  194. #[test]
  195. fn extracts_terminal_slash_command_prefixes_with_arguments() {
  196. assert_eq!(slash_command_prefix("/he", 3), Some("/he"));
  197. assert_eq!(slash_command_prefix("/help me", 8), Some("/help me"));
  198. assert_eq!(
  199. slash_command_prefix("/session switch ses", 19),
  200. Some("/session switch ses")
  201. );
  202. assert_eq!(slash_command_prefix("hello", 5), None);
  203. assert_eq!(slash_command_prefix("/help", 2), None);
  204. }
  205. #[test]
  206. fn completes_matching_slash_commands() {
  207. let helper = SlashCommandHelper::new(vec![
  208. "/help".to_string(),
  209. "/hello".to_string(),
  210. "/status".to_string(),
  211. ]);
  212. let history = DefaultHistory::new();
  213. let ctx = Context::new(&history);
  214. let (start, matches) = helper
  215. .complete("/he", 3, &ctx)
  216. .expect("completion should work");
  217. assert_eq!(start, 0);
  218. assert_eq!(
  219. matches
  220. .into_iter()
  221. .map(|candidate| candidate.replacement)
  222. .collect::<Vec<_>>(),
  223. vec!["/help".to_string(), "/hello".to_string()]
  224. );
  225. }
  226. #[test]
  227. fn completes_matching_slash_command_arguments() {
  228. let helper = SlashCommandHelper::new(vec![
  229. "/model".to_string(),
  230. "/model opus".to_string(),
  231. "/model sonnet".to_string(),
  232. "/session switch alpha".to_string(),
  233. ]);
  234. let history = DefaultHistory::new();
  235. let ctx = Context::new(&history);
  236. let (start, matches) = helper
  237. .complete("/model o", 8, &ctx)
  238. .expect("completion should work");
  239. assert_eq!(start, 0);
  240. assert_eq!(
  241. matches
  242. .into_iter()
  243. .map(|candidate| candidate.replacement)
  244. .collect::<Vec<_>>(),
  245. vec!["/model opus".to_string()]
  246. );
  247. }
  248. #[test]
  249. fn ignores_non_slash_command_completion_requests() {
  250. let helper = SlashCommandHelper::new(vec!["/help".to_string()]);
  251. let history = DefaultHistory::new();
  252. let ctx = Context::new(&history);
  253. let (_, matches) = helper
  254. .complete("hello", 5, &ctx)
  255. .expect("completion should work");
  256. assert!(matches.is_empty());
  257. }
  258. #[test]
  259. fn tracks_current_buffer_through_highlighter() {
  260. let helper = SlashCommandHelper::new(Vec::new());
  261. let _ = helper.highlight("draft", 5);
  262. assert_eq!(helper.current_line(), "draft");
  263. }
  264. #[test]
  265. fn push_history_ignores_blank_entries() {
  266. let mut editor = LineEditor::new("> ", vec!["/help".to_string()]);
  267. editor.push_history(" ");
  268. editor.push_history("/help");
  269. assert_eq!(editor.editor.history().len(), 1);
  270. }
  271. #[test]
  272. fn set_completions_replaces_and_normalizes_candidates() {
  273. let mut editor = LineEditor::new("> ", vec!["/help".to_string()]);
  274. editor.set_completions(vec![
  275. "/model opus".to_string(),
  276. "/model opus".to_string(),
  277. "status".to_string(),
  278. ]);
  279. let helper = editor.editor.helper().expect("helper should exist");
  280. assert_eq!(helper.completions, vec!["/model opus".to_string()]);
  281. }
  282. }