use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread::JoinHandle;
use std::time::Duration;
use bevy_ecs::prelude::*;
use bevy_ecs::schedule::ScheduleLabel;
use bevy_ecs::system::SystemParam;
use bevy_log::error;
use crossbeam_channel::Receiver;
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyEventKind};
use ratatui::layout::Rect;
use super::clipboard;
use super::complete::{self, Completion, CompletionKind, CompletionSource};
use super::editor::{Editor, MainEditor};
use super::history::{PromptHistory, Sent};
use super::keys::{self, Found, InputAction, KeyBinding, KeyFocus, KeyPress, KeyPressed, KeyScope};
use super::pane::{MainPane, PaneState};
use super::panel::{TerminalSize, TuiPanel};
use super::pick::{self, Choices, PickerInput, Showing};
use super::render::selectable;
use super::terminal::TerminalFocus;
use super::view::{ItemRow, TranscriptItem, TranscriptView, TuiView};
use rig_ecs::agent::{ActiveTurn, Notice, StopTurn};
use rig_ecs::inbox::{Deliver, DeliveryMode};
use rig_ecs::tasks::Wake;
use rig_harness::main_agent::ShownAgent;
use rig_harness::pick::{Chosen, Pick};
use rig_harness::prelude::{Processes, Scratch};
const PAGE: usize = 10;
const POLL: Duration = Duration::from_millis(100);
const KEYS_PER_FRAME: usize = 1024;
#[derive(Resource, Clone, Debug, Default, PartialEq, Eq)]
pub struct PendingInput(pub Vec<KeyInput>);
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum KeyInput {
Event(Event),
Action(InputAction),
}
#[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
pub struct TuiKeys;
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum InputSystems {
Filter,
Dispatch,
}
#[derive(Resource, Clone, Debug, Default, PartialEq, Eq)]
pub struct InputContext {
scope: KeyScope,
editor: Option<Entity>,
completion: Option<Completion>,
typed: Option<(CompletionKind, String)>,
answering: Option<Entity>,
picking: Option<Entity>,
selection: Option<usize>,
}
impl InputContext {
pub fn scope(&self) -> KeyScope {
self.scope
}
pub fn editor(&self) -> Option<Entity> {
self.editor
}
pub fn completion(&self) -> Option<&Completion> {
self.completion.as_ref()
}
pub fn typed(&self) -> Option<(CompletionKind, &str)> {
let (kind, text) = self.typed.as_ref()?;
Some((*kind, text))
}
pub fn answering(&self) -> Option<Entity> {
self.answering
}
pub fn picking(&self) -> Option<Entity> {
self.picking
}
pub fn selection(&self) -> Option<usize> {
self.selection
}
}
#[derive(Resource)]
pub(crate) struct TerminalInput {
events: Receiver<Event>,
stop: Arc<AtomicBool>,
thread: Option<JoinHandle<()>>,
}
impl TerminalInput {
pub(crate) fn start(wake: Wake) -> std::io::Result<Self> {
let (sender, events) = crossbeam_channel::unbounded();
let stop = Arc::new(AtomicBool::new(false));
let shared = Arc::clone(&stop);
let thread = std::thread::Builder::new()
.name("rig-harness-input".to_owned())
.spawn(move || {
while !shared.load(Ordering::Relaxed) {
let event = match event::poll(POLL) {
Ok(false) => continue,
Ok(true) => event::read(),
Err(failure) => Err(failure),
};
match event {
Ok(event) => {
if sender.send(event).is_err() {
return;
}
wake.wake();
}
Err(failure) => {
error!("could not read the terminal: {failure}");
return;
}
}
}
})?;
Ok(Self {
events,
stop,
thread: Some(thread),
})
}
pub(crate) fn stop(mut self) {
self.stop.store(true, Ordering::Relaxed);
if let Some(thread) = self.thread.take() {
thread.join().ok();
}
}
}
impl Drop for TerminalInput {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
}
}
pub(crate) fn receive_events(input: Res<TerminalInput>, mut events: ResMut<PendingInput>) {
let mut read = input.events.try_iter().peekable();
if read.peek().is_some() {
events.0.extend(read.map(KeyInput::Event));
}
}
pub(crate) fn clear_events(mut events: ResMut<PendingInput>) {
if !events.0.is_empty() {
events.0.clear();
}
}
pub(crate) fn run_keys(world: &mut World) {
if world
.get_resource::<PendingInput>()
.is_none_or(|events| events.0.is_empty())
{
return;
}
let mut events = world.resource_mut::<PendingInput>();
let mut queue: VecDeque<KeyInput> = std::mem::take(&mut events.0).into();
let mut done = 0;
while let Some(item) = queue.pop_front() {
if done == KEYS_PER_FRAME {
queue.push_front(item);
world.resource_mut::<PendingInput>().0 = queue.into();
if let Some(wake) = world.get_resource::<Wake>() {
wake.wake();
}
return;
}
done += 1;
world.resource_mut::<PendingInput>().0.push(item);
let text = main_text(world);
world.run_schedule(TuiKeys);
let typed = main_text(world) != text;
if typed {
let mut main = world.query_filtered::<&mut TranscriptView, With<MainPane>>();
for mut transcript in main.iter_mut(world) {
if transcript.selected.is_some() {
transcript.selected = None;
}
}
}
let left = std::mem::take(&mut world.resource_mut::<PendingInput>().0);
for item in left.into_iter().rev() {
queue.push_front(item);
}
}
}
fn main_text(world: &mut World) -> Option<String> {
let mut main = world.query_filtered::<&Editor, With<MainEditor>>();
let text = main.single(world).ok().map(Editor::text);
text.map(str::to_owned)
}
pub(crate) fn note_input_context(
(view, main): (Res<TuiView>, Query<&TranscriptView, With<MainPane>>),
showing: Query<(Entity, &PickerInput)>,
focus: Res<KeyFocus>,
panels: Query<(), With<TuiPanel>>,
choices: Choices,
(editors, sources): (
Query<(Entity, &Editor, Has<MainEditor>)>,
Query<&CompletionSource>,
),
mut state: ResMut<InputContext>,
) {
let main_editor = editors.iter().find(|(.., main)| *main);
let typed = main_editor.and_then(|(_, input, _)| complete::typed(input, &sources));
let focused = focus.0.filter(|panel| editors.contains(*panel));
let editor = focused.or(main_editor.map(|(entity, ..)| entity));
let scope = match focus.0 {
Some(panel) if panels.contains(panel) => KeyScope::Panel(panel),
_ if reading(&main) => KeyScope::Transcript,
_ => KeyScope::Input,
};
let choosing = showing.iter().next();
let answering = choosing.filter(|(_, c)| c.draft.is_some()).map(|(e, _)| e);
let picker = choosing.filter(|(_, c)| c.draft.is_none());
let picking = picker.map(|(e, _)| e);
let selection = picker.and_then(|(entity, choosing)| {
let shown = choices.shown(entity)?;
match shown.rows(&choosing.filter).get(choosing.selected) {
Some(&pick::Row::Item(index)) => Some(index),
_ => None,
}
});
state.set_if_neq(InputContext {
scope,
editor,
completion: view.completion.clone(),
typed,
answering,
picking,
selection,
});
}
fn reading(main: &Query<&TranscriptView, With<MainPane>>) -> bool {
main.iter().any(|view| view.selected.is_some())
}
#[derive(SystemParam)]
pub(crate) struct Keymap<'w, 's> {
bindings: Query<'w, 's, (Entity, &'static KeyBinding, Option<&'static InputAction>)>,
focus: ResMut<'w, KeyFocus>,
panels: Query<'w, 's, (), With<TuiPanel>>,
}
impl Keymap<'_, '_> {
fn scope(&mut self, reading: bool) -> KeyScope {
let unfocused = if reading {
KeyScope::Transcript
} else {
KeyScope::Input
};
match self.focus.0 {
Some(panel) if self.panels.contains(panel) => KeyScope::Panel(panel),
Some(_) => {
self.focus.0 = None;
unfocused
}
None => unfocused,
}
}
}
pub(crate) struct Context<'a, 'w, 's> {
agent: Option<Entity>,
busy: bool,
scratch: Option<&'a Scratch>,
processes: Option<&'a Processes>,
wake: &'a Wake,
history: &'a [String],
choices: &'a Choices<'w, 's>,
}
pub(crate) struct Target<'a> {
view: &'a mut TuiView,
input: &'a mut Editor,
pane: Option<(Mut<'a, TranscriptView>, Mut<'a, PaneState>)>,
}
type Panes<'w, 's> = Query<
'w,
's,
(
Entity,
&'static mut TranscriptView,
&'static mut PaneState,
Has<MainPane>,
),
>;
type Editors<'w, 's> = Query<'w, 's, (Entity, &'static mut Editor, Has<MainEditor>)>;
fn keyed_editor(editors: &Editors, focus: Option<Entity>) -> Option<Entity> {
let focused = focus.filter(|entity| editors.contains(*entity));
let main = || editors.iter().find(|(.., main)| *main);
focused.or_else(|| main().map(|(entity, ..)| entity))
}
fn keyed_pane(panes: &Panes, focus: Option<Entity>) -> Option<Entity> {
let focused = focus.filter(|entity| panes.contains(*entity));
let main = || {
panes
.iter()
.find(|(.., main)| *main)
.map(|(entity, ..)| entity)
};
focused.or_else(main)
}
pub(crate) fn read_input(
mut events: ResMut<PendingInput>,
(mut view, mut panes, shown): (ResMut<TuiView>, Panes, Res<ShownAgent>),
(mut editors, mut showing): (Editors, Showing),
mut keymap: Keymap,
choices: Choices,
agents: Query<Has<ActiveTurn>>,
(scratch, processes, wake, history): (
Option<Res<Scratch>>,
Option<Res<Processes>>,
Res<Wake>,
Res<PromptHistory>,
),
(mut screen, mut focus): (ResMut<TerminalSize>, ResMut<TerminalFocus>),
mut commands: Commands,
) {
if events.0.is_empty() {
return;
}
let busy = shown
.0
.and_then(|agent| agents.get(agent).ok())
.unwrap_or(false);
let context = Context {
agent: shown.0,
busy,
scratch: scratch.as_deref(),
processes: processes.as_deref(),
wake: &wake,
history: &history.entries,
choices: &choices,
};
let mut none = Editor::default();
for item in std::mem::take(&mut events.0) {
let editor = keyed_editor(&editors, keymap.focus.0);
let in_panel = editor.is_some() && editor == keymap.focus.0;
let mut found = editor.and_then(|editor| editors.get_mut(editor).ok());
let input: &mut Editor = match &mut found {
Some((_, editor, _)) => editor,
None => &mut none,
};
let pane = keyed_pane(&panes, keymap.focus.0);
let reading = panes
.iter()
.any(|(_, view, _, main)| main && view.selected.is_some());
let event = match item {
KeyInput::Event(event) => event,
KeyInput::Action(action) => {
let pane = pane.and_then(|pane| panes.get_mut(pane).ok());
let mut target = Target {
view: &mut view,
input,
pane: pane.map(|(_, view, state, _)| (view, state)),
};
act(action, &mut target, (&context, &showing), &mut commands);
continue;
}
};
match event {
Event::Key(key) if key.kind == KeyEventKind::Release => {
released(key, (shown.0, reading), &mut keymap, &mut commands);
}
Event::Key(key) => {
let picking = showing.iter_mut().find(|(_, c)| c.draft.is_none());
if let Some((entity, mut choosing)) = picking {
let shown = choices.shown(entity);
let chosen =
shown.and_then(|shown| pick::choice_key(key, &mut choosing, &shown));
if let Some(chosen) = chosen {
commands.trigger(Pick::person(entity, chosen));
view.queue.retain(|&queued| queued != entity);
pick::show_first(&view, &choices, (&showing, input), &mut commands);
}
continue;
}
let pane = pane.and_then(|pane| panes.get_mut(pane).ok());
let mut target = Target {
view: &mut view,
input,
pane: pane.map(|(_, view, state, _)| (view, state)),
};
pressed(
key,
(reading, in_panel),
&mut target,
&mut keymap,
(&context, &showing),
&mut commands,
);
}
Event::Paste(text) => match showing.iter_mut().find(|(_, c)| c.draft.is_none()) {
Some((_, mut choosing)) => {
choosing
.filter
.push_str(text.lines().next().unwrap_or_default());
choosing.selected = 0;
}
None if !in_panel && matches!(keymap.scope(reading), KeyScope::Panel(_)) => {}
None => match clipboard::dropped_image(&text) {
Some(path) => clipboard::type_path(input, &path),
None => input.insert(&text),
},
},
Event::Resize(width, height) => {
screen.set_if_neq(TerminalSize(Rect::new(0, 0, width, height)));
}
Event::FocusGained => {
focus.set_if_neq(TerminalFocus(Some(true)));
}
Event::FocusLost => {
focus.set_if_neq(TerminalFocus(Some(false)));
}
Event::Mouse(_) => {}
}
}
}
fn pressed(
key: KeyEvent,
(reading, in_panel): (bool, bool),
target: &mut Target,
keymap: &mut Keymap,
(context, showing): (&Context, &Showing),
commands: &mut Commands,
) {
let Target { view, input, .. } = target;
let press = KeyPress::of(&key);
let scope = keymap.scope(reading);
let bindings: Vec<(Entity, &KeyBinding)> = keymap
.bindings
.iter()
.map(|(entity, binding, _)| (entity, binding))
.collect();
view.pending.push(press);
let found = keys::find(&bindings, &view.pending, scope);
if found == Found::Partial {
return;
}
let sequence = std::mem::take(&mut view.pending).len() > 1;
let Found::Binding(entity) = found else {
if sequence {
return;
}
match scope {
KeyScope::Panel(_) if press.code == KeyCode::Esc => keymap.focus.0 = None,
KeyScope::Panel(_) if !in_panel => {}
_ => {
if let Some(character) = press.typed() {
input.insert_char(character);
}
}
}
return;
};
let Ok((_, binding, action)) = keymap.bindings.get(entity) else {
return;
};
if binding.releases && key.kind == KeyEventKind::Repeat {
return;
}
match action {
Some(action) => act(*action, target, (context, showing), commands),
None => commands.trigger(KeyPressed {
entity,
agent: context.agent,
released: false,
}),
}
}
fn released(
key: KeyEvent,
(agent, reading): (Option<Entity>, bool),
keymap: &mut Keymap,
commands: &mut Commands,
) {
let press = KeyPress::of(&key);
let scope = keymap.scope(reading);
let released = keymap.bindings.iter().find(|(_, binding, _)| {
binding.releases
&& binding.keys.presses() == [press]
&& scope.applying().any(|applying| applying == binding.scope)
});
if let Some((entity, ..)) = released {
commands.trigger(KeyPressed {
entity,
agent,
released: true,
});
}
}
fn act(
action: InputAction,
target: &mut Target,
(context, showing): (&Context, &Showing),
commands: &mut Commands,
) {
use InputAction as Do;
let Target { view, input, pane } = target;
let editing = showing
.iter()
.find(|(_, choosing)| choosing.draft.is_some());
let editing = editing.map(|(entity, _)| entity);
match (action, editing) {
(Do::Send, Some(ask)) if view.completion.as_ref().is_none_or(|c| c.is_exact()) => {
let answer = Chosen::text(input.take());
return close(ask, Some(answer), target, (context, showing), commands);
}
(Do::Stop, Some(ask)) if view.completion.is_none() => {
return close(ask, None, target, (context, showing), commands);
}
_ => {}
}
match action {
Do::Clear => input.clear(),
Do::NewLine => input.insert_char('\n'),
Do::LineStart => input.home(),
Do::LineEnd => input.end(),
Do::Left => input.left(),
Do::Right => input.right(),
Do::WordLeft => input.word_left(),
Do::WordRight => input.word_right(),
Do::Backspace => input.backspace(),
Do::Delete => input.delete(),
Do::DeleteWordBefore => input.delete_word_before(),
Do::DeleteWordAfter => input.delete_word_after(),
Do::DeleteToLineStart => input.delete_to_line_start(),
Do::DeleteToLineEnd => input.delete_to_line_end(),
Do::LineUp => input.up(context.history),
Do::LineDown => input.down(context.history),
Do::Up => match &mut view.completion {
Some(completion) => completion.select(-1),
None => input.up(context.history),
},
Do::Down => match &mut view.completion {
Some(completion) => completion.select(1),
None => input.down(context.history),
},
Do::ScrollUp | Do::ScrollDown | Do::PageUp | Do::PageDown => {
if let Some((_, state)) = pane {
let page = PAGE.min(state.page());
match action {
Do::ScrollUp => state.scroll.up(1),
Do::ScrollDown => state.scroll.down(1),
Do::PageUp => state.scroll.up(page),
_ => state.scroll.down(page),
}
}
}
Do::Tab if view.completion.is_none() && context.busy && editing.is_none() => {
let state = pane.as_mut().map(|(_, state)| &mut **state);
send(view, input, state, context.agent, commands, true);
}
Do::Tab => accept_completion(view, input),
Do::Send => {
if view
.completion
.as_ref()
.is_some_and(|completion| !completion.is_exact())
{
accept_completion(view, input);
return;
}
if input.continue_line() {
return;
}
let state = pane.as_mut().map(|(_, state)| &mut **state);
send(view, input, state, context.agent, commands, false);
}
Do::Stop => match view.completion.take() {
Some(completion) => view.dismissed = Some(completion.start),
None => {
if let Some(entity) = context.agent {
commands.trigger(StopTurn::new(entity));
}
}
},
Do::PasteImage => {
let Some(images) = context.scratch.and_then(Scratch::images) else {
let why = "This session keeps no files to paste an image into.";
commands.write_message(Notice::info(context.agent, why));
return;
};
let launcher = context
.processes
.map(Processes::launcher)
.unwrap_or_default();
clipboard::paste_image(images, launcher, context.wake, commands);
}
Do::SelectUp | Do::SelectDown => {
if let Some((transcript, state)) = pane {
select(action == Do::SelectUp, transcript, state);
}
}
Do::Deselect => {
if let Some((transcript, _)) = pane
&& transcript.selected.is_some()
{
transcript.selected = None;
}
}
Do::AnswerInline(number) => {
answer_inline(usize::from(number), target, (context, showing), commands);
}
Do::DismissInline => {
if let Some(&ask) = view.inline.first() {
view.inline.remove(0);
commands.trigger(Pick::person(ask, None));
}
}
}
}
fn close(
ask: Entity,
chosen: Option<Chosen>,
target: &mut Target,
(context, showing): (&Context, &Showing),
commands: &mut Commands,
) {
commands.trigger(Pick::person(ask, chosen));
target.view.queue.retain(|&queued| queued != ask);
pick::show_first(
target.view,
context.choices,
(showing, target.input),
commands,
);
}
fn select(up: bool, transcript: &mut Mut<TranscriptView>, state: &mut Mut<PaneState>) {
let items = state
.transcript
.marks_in(&state.parts)
.map(|(item, _)| item);
let items: Vec<&TranscriptItem> = items.filter(|item| selectable(item)).collect();
let at = (transcript.selected.as_ref())
.and_then(|selected| items.iter().position(|item| *item == selected));
let next = match (at, up) {
(None, true) => items.len().checked_sub(1),
(None, false) => None,
(Some(at), true) => Some(at.saturating_sub(1)),
(Some(at), false) => Some(at + 1),
};
let selected = next
.and_then(|next| items.get(next))
.map(|item| (*item).clone());
if let Some(item) = &selected {
state.scroll.to(item.clone(), ItemRow::Reveal);
}
if transcript.selected != selected {
transcript.selected = selected;
}
}
fn answer_inline(
number: usize,
target: &mut Target,
(context, showing): (&Context, &Showing),
commands: &mut Commands,
) {
let Some(&ask) = target.view.inline.first() else {
return;
};
let Some(shown) = context.choices.shown(ask) else {
return;
};
let Some(index) = number.checked_sub(1) else {
return;
};
if index < shown.items.len() {
target.view.inline.remove(0);
commands.trigger(Pick::person(ask, Some(Chosen::item(index))));
} else if index == shown.items.len() && shown.free_text && showing.is_empty() {
target.view.inline.remove(0);
target.view.queue.insert(0, ask);
pick::show_first(
target.view,
context.choices,
(showing, target.input),
commands,
);
}
}
pub(crate) fn send(
view: &mut TuiView,
input: &mut Editor,
pane: Option<&mut PaneState>,
agent: Option<Entity>,
commands: &mut Commands,
queue: bool,
) {
let Some(entity) = agent else {
return;
};
if input.text().trim().is_empty() {
return;
}
let text = input.take();
commands.write_message(Sent(text.clone()));
if let Some(pane) = pane {
pane.scroll.follow();
}
let mode = if queue {
DeliveryMode::Queue
} else {
DeliveryMode::Steer
};
let refused = view.refused.take().as_ref() == Some(&text);
let typed = Deliver::new(entity, text, mode);
commands.trigger(if refused { typed.literal() } else { typed });
}
fn accept_completion(view: &mut TuiView, input: &mut Editor) {
let Some(completion) = view.completion.take() else {
return;
};
if let Some(text) = completion.chosen() {
input.replace_before_cursor(completion.start, &text);
}
}
#[cfg(test)]
mod tests;