From 6ea760bb6b36da241961e8ecd60cb2c5e15c0a78 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Tue, 24 Feb 2026 11:48:20 -0800 Subject: feat: Generate commands or ask questions with `atuin ai` (#3199) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR refines the system created in #3178 to be suitable for a v1 release. --- ## Overview `atuin-ai` is a separate binary that allows for generating commands and asking questions from the command line. It is fully opt-in. ## Usage `atuin ai init` will output bindings for your shell. Currently, bash, zsh, and fish are supported. ```bash eval "$(atuin ai init)" ``` Once the hooks are installed, just press `?` on an empty prompt line to call up the TUI. `atuin ai` requires an account on [Atuin Hub](https://hub.atuin.sh/); you will be prompted to log in on first use. ## Features ### Command generation Prompt the LLM to create a command, and get one back, no fuss. Press `enter` to run, or `tab` to insert. ``` ┌Ask questions or generate a command:──────────────────────────┐ │ │ │ > Get a list of running docker containers │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ docker ps │ │ │ └────[Enter]: Run [Tab]: Insert [f]: Follow-up [Esc]: Cancel┘ ``` ### Follow-up You can follow-up with `f` to specify a refinement prompt to update the command that will be inserted. ``` ┌Ask questions or generate a command:──────────────────────────┐ │ │ │ > Get a list of running docker containers │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ docker ps │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ > Actually I want to get all docker containers │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ docker ps -a │ │ │ └────[Enter]: Run [Tab]: Insert [f]: Follow-up [Esc]: Cancel┘ ``` You can also follow-up with questions to get responses in natural language. ``` ┌Ask questions or generate a command:──────────────────────────┐ │ │ │ > Get a list of running docker containers │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ docker ps │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ > Actually I want to get all docker containers │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ docker ps -a │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ > What other useful flags to `docker ps` should I know? │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ Here are some handy `docker ps` flags: │ │ │ │ - `-q` — Only show container IDs (great for piping to │ │ other commands) │ │ - `-s` — Show container sizes │ │ - `-n 5` — Show the last 5 created containers │ │ - `-l` — Show only the latest created container │ │ - `--no-trunc` — Don't truncate output (shows full IDs and │ │ commands) │ │ - `-f` or `--filter` — Filter by condition, e.g.: │ │ - `-f status=exited` — only exited containers │ │ - `-f name=myapp` — filter by name │ │ - `-f ancestor=nginx` — filter by image │ │ - `--format` — Custom output using Go templates, e.g.: │ │ `--format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"` │ │ │ │ A common combo is `docker ps -aq` to get all container │ │ IDs, useful for bulk operations like `docker rm $(docker │ │ ps -aq)`. │ │ │ └────[Enter]: Run [Tab]: Insert [f]: Follow-up [Esc]: Cancel┘ ``` You can use `enter` or `tab` at any time to run or insert the last suggested command, even if it was suggested in a previous turn. ### Conversational and search usage If you prompt the LLM with a question that doesn't imply you want to generate a command, it can respond in natural language, and use web search if necessary to fetch the data it needs. ``` ┌Ask questions or generate a command:──────────────────────────┐ │ │ │ > What is the latest version of atuin? │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ ✓ Used 2 tools │ │ │ │ The latest version of Atuin is **v18.12.0**, available on │ │ the [GitHub releases │ │ page](https://github.com/atuinsh/atuin/releases). │ │ │ └─────────────────────────────────[f]: Follow-up [Esc]: Cancel┘ ``` ### Dangerous or low-confidence command detection The LLM scores its confidence in the command, as well as how dangerous the command is. This information is shown if a threshold is exceeded, and requires an extra confirmation step before running automatically with `enter`. The Atuin Hub server also monitors suggested commands for dangerous patterns the LLM didn't catch, and appends its own assessment at the end of the LLM's own assessment. ``` ┌Ask questions or generate a command:──────────────────────────┐ │ │ │ > Delete all files from $HOME │ │ │ ├──────────────────────────────────────────────────────────────┤ │ │ │ $ rm -rf $HOME/* │ │ │ │ ! ⚠️ This will PERMANENTLY delete ALL files and directories │ │ in your home directory, including documents, downloads, │ │ configurations, SSH keys, and everything else. This is │ │ irreversible and will likely break your system. Also note │ │ this won't delete hidden (dot) files — if you want those │ │ too, that's even more destructive.; [Server] Recursive │ │ delete of critical directory │ │ │ └────[Enter]: Run [Tab]: Insert [f]: Follow-up [Esc]: Cancel┘ ``` --------- Co-authored-by: Claude Opus 4.5 --- crates/atuin-ai/Cargo.toml | 13 +- crates/atuin-ai/render-tests.sh | 34 + crates/atuin-ai/replay-states.sh | 101 +++ crates/atuin-ai/src/commands.rs | 71 +- crates/atuin-ai/src/commands/debug_render.rs | 460 +++++++++++++ crates/atuin-ai/src/commands/init.rs | 155 ++++- crates/atuin-ai/src/commands/inline.rs | 932 ++++++++++++++------------- crates/atuin-ai/src/main.rs | 1 + crates/atuin-ai/src/tui/app.rs | 157 +++++ crates/atuin-ai/src/tui/event.rs | 303 +++++++++ crates/atuin-ai/src/tui/mod.rs | 14 + crates/atuin-ai/src/tui/render.rs | 674 +++++++++++++++++++ crates/atuin-ai/src/tui/spinner.rs | 99 +++ crates/atuin-ai/src/tui/state.rs | 530 +++++++++++++++ crates/atuin-ai/src/tui/terminal.rs | 203 ++++++ crates/atuin-ai/src/tui/view_model.rs | 400 ++++++++++++ crates/atuin-ai/test-renders.json | 295 +++++++++ crates/atuin-client/src/settings.rs | 17 + 18 files changed, 3984 insertions(+), 475 deletions(-) create mode 100755 crates/atuin-ai/render-tests.sh create mode 100755 crates/atuin-ai/replay-states.sh create mode 100644 crates/atuin-ai/src/commands/debug_render.rs create mode 100644 crates/atuin-ai/src/tui/app.rs create mode 100644 crates/atuin-ai/src/tui/event.rs create mode 100644 crates/atuin-ai/src/tui/mod.rs create mode 100644 crates/atuin-ai/src/tui/render.rs create mode 100644 crates/atuin-ai/src/tui/spinner.rs create mode 100644 crates/atuin-ai/src/tui/state.rs create mode 100644 crates/atuin-ai/src/tui/terminal.rs create mode 100644 crates/atuin-ai/src/tui/view_model.rs create mode 100644 crates/atuin-ai/test-renders.json (limited to 'crates') diff --git a/crates/atuin-ai/Cargo.toml b/crates/atuin-ai/Cargo.toml index d42b6a67..692ff9c2 100644 --- a/crates/atuin-ai/Cargo.toml +++ b/crates/atuin-ai/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "atuin-ai" edition = "2024" -description = "ai library for atuin" +description = "AI integration for Atuin CLI" rust-version = { workspace = true } version = { workspace = true } @@ -30,8 +30,15 @@ tracing-appender = "0.2.4" reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -crossterm = { workspace = true, features = ["use-dev-tty"] } -ratatui = { workspace = true } +crossterm = { workspace = true, features = ["use-dev-tty", "event-stream"] } +ratatui = { workspace = true, features = ["unstable-rendered-line-info"] } +futures = "0.3" +eventsource-stream = "0.2" +pulldown-cmark = "0.13.0" +async-stream = "0.3" +uuid = { workspace = true } +tui-textarea-2 = "0.9.1" +unicode-width = "0.2" [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/crates/atuin-ai/render-tests.sh b/crates/atuin-ai/render-tests.sh new file mode 100755 index 00000000..8dedc76e --- /dev/null +++ b/crates/atuin-ai/render-tests.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Render all test cases from test-renders.json +# Usage: ./render-tests.sh [test_name] +# With no args: renders all tests +# With arg: renders only matching test (e.g., ./render-tests.sh 05) + +set -e +cd "$(dirname "$0")" + +JSON_FILE="test-renders.json" +FILTER="${1:-}" + +# Build once +cargo build -p atuin-ai --quiet + +# Count tests +TOTAL=$(jq length "$JSON_FILE") + +for i in $(seq 0 $((TOTAL - 1))); do + NAME=$(jq -r ".[$i].name" "$JSON_FILE") + DESC=$(jq -r ".[$i].description" "$JSON_FILE") + STATE=$(jq -c ".[$i].state" "$JSON_FILE") + + # Skip if filter provided and doesn't match + if [[ -n "$FILTER" && ! "$NAME" =~ $FILTER ]]; then + continue + fi + + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$NAME] $DESC" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain + echo "" +done diff --git a/crates/atuin-ai/replay-states.sh b/crates/atuin-ai/replay-states.sh new file mode 100755 index 00000000..4f586709 --- /dev/null +++ b/crates/atuin-ai/replay-states.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# Replay state snapshots from a debug state JSONL file +# Usage: ./replay-states.sh [entry-number] +# With no entry: renders all frames in sequence (press Enter to advance) +# With entry number: renders just that frame + +set -e +# cd "$(dirname "$0")" + +STATE_FILE="${1:-}" +ENTRY_FILTER="${2:-}" + +if [[ -z "$STATE_FILE" ]]; then + echo "Usage: $0 [entry-number]" + echo "" + echo "Examples:" + echo " $0 /tmp/state.jsonl # Interactive replay of all frames" + echo " $0 /tmp/state.jsonl 15 # Show just entry 15" + exit 1 +fi + +if [[ ! -f "$STATE_FILE" ]]; then + echo "Error: File not found: $STATE_FILE" + exit 1 +fi + +# Build once +cargo build -p atuin-ai --quiet + +# Count entries +TOTAL=$(wc -l < "$STATE_FILE" | tr -d ' ') + +if [[ -n "$ENTRY_FILTER" ]]; then + # Show single entry + LINE=$(sed -n "${ENTRY_FILTER}p" "$STATE_FILE") + if [[ -z "$LINE" ]]; then + echo "Error: Entry $ENTRY_FILTER not found (file has $TOTAL entries)" + exit 1 + fi + + ENTRY=$(echo "$LINE" | jq -r '.entry') + LABEL=$(echo "$LINE" | jq -r '.label') + STATE=$(echo "$LINE" | jq -c '.state') + + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$ENTRY/$TOTAL] $LABEL" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain +else + # Interactive replay + echo "Replaying $TOTAL frames from $STATE_FILE" + echo "Press Enter to advance, 'q' to quit, or number+Enter to jump" + echo "" + + CURRENT=1 + while [[ $CURRENT -le $TOTAL ]]; do + LINE=$(sed -n "${CURRENT}p" "$STATE_FILE") + ENTRY=$(echo "$LINE" | jq -r '.entry') + LABEL=$(echo "$LINE" | jq -r '.label') + STATE=$(echo "$LINE" | jq -c '.state') + + clear + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$CURRENT/$TOTAL] $LABEL" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain + echo "" + echo "[Enter: next] [p: prev] [number: jump] [s: show state JSON] [q: quit]" + + read -r INPUT + case "$INPUT" in + q|Q) + break + ;; + p|P) + if [[ $CURRENT -gt 1 ]]; then + CURRENT=$((CURRENT - 1)) + fi + ;; + s|S) + echo "" + echo "State JSON:" + echo "$STATE" | jq . + echo "" + echo "Press Enter to continue..." + read -r + ;; + ''|' ') + CURRENT=$((CURRENT + 1)) + ;; + *[0-9]*) + if [[ "$INPUT" =~ ^[0-9]+$ ]] && [[ "$INPUT" -ge 1 ]] && [[ "$INPUT" -le $TOTAL ]]; then + CURRENT=$INPUT + else + echo "Invalid entry number (1-$TOTAL)" + sleep 1 + fi + ;; + esac + done +fi diff --git a/crates/atuin-ai/src/commands.rs b/crates/atuin-ai/src/commands.rs index 56741544..7d5ca16b 100644 --- a/crates/atuin-ai/src/commands.rs +++ b/crates/atuin-ai/src/commands.rs @@ -1,7 +1,11 @@ +use atuin_common::shell::Shell; use clap::{Parser, Subcommand}; use tracing::Level; use tracing_subscriber::{EnvFilter, Layer, fmt, layer::SubscriberExt, util::SubscriberInitExt}; +#[cfg(debug_assertions)] +pub mod debug_render; + pub mod init; pub mod inline; @@ -16,6 +20,10 @@ struct Cli { #[arg(long, global = true, env = "ATUIN_AI_API_ENDPOINT")] api_endpoint: Option, + /// Custom API token + #[arg(long, global = true, env = "ATUIN_AI_API_TOKEN")] + api_token: Option, + #[command(subcommand)] command: Commands, } @@ -23,13 +31,10 @@ struct Cli { #[derive(Subcommand, Debug)] enum Commands { /// Initialize shell integration - Init, - - /// Complete current command line - Complete { - /// Current command line to complete - #[arg(value_name = "COMMAND")] - command: Option, + Init { + /// Shell to generate integration for; defaults to "auto" + #[arg(value_name = "SHELL", default_value = "auto")] + shell: String, }, /// Inline completion mode with small TUI overlay @@ -41,10 +46,27 @@ enum Commands { /// Start in natural language mode #[arg(long)] natural_language: bool, + + /// Keep TUI output visible after exit (default: erase) + #[arg(long)] + keep: bool, + + /// Log state changes to file for debugging (dev tool) + #[arg(long, value_name = "FILE")] + debug_state: Option, }, - /// Interactive mode with TUI - Interactive, + /// Debug render: output a single frame from JSON state (dev tool) + #[cfg(debug_assertions)] + DebugRender { + /// Input file (reads from stdin if not provided) + #[arg(short, long)] + input: Option, + + /// Output format: ansi (default), plain, json + #[arg(short, long, default_value = "ansi")] + format: String, + }, } pub async fn run() -> eyre::Result<()> { @@ -53,13 +75,32 @@ pub async fn run() -> eyre::Result<()> { init_tracing(cli.verbose); match cli.command { - Commands::Init => init::run().await, + Commands::Init { shell } => init::run(shell).await, Commands::Inline { command, natural_language, - } => inline::run(command, natural_language, cli.api_endpoint).await, - Commands::Complete { command } => inline::run(command, false, cli.api_endpoint).await, - Commands::Interactive => Err(eyre::eyre!("interactive mode not implemented yet")), + keep, + debug_state, + } => { + inline::run( + command, + natural_language, + cli.api_endpoint, + cli.api_token, + keep, + debug_state, + ) + .await + } + #[cfg(debug_assertions)] + Commands::DebugRender { input, format } => { + let output_format = match format.as_str() { + "plain" => debug_render::OutputFormat::Plain, + "json" => debug_render::OutputFormat::Json, + _ => debug_render::OutputFormat::Ansi, + }; + debug_render::run(input, output_format).await + } } } @@ -95,3 +136,7 @@ fn init_tracing(verbose: bool) { subscriber.init(); } } + +pub fn detect_shell() -> Option { + Some(Shell::current().to_string()) +} diff --git a/crates/atuin-ai/src/commands/debug_render.rs b/crates/atuin-ai/src/commands/debug_render.rs new file mode 100644 index 00000000..e78a418a --- /dev/null +++ b/crates/atuin-ai/src/commands/debug_render.rs @@ -0,0 +1,460 @@ +//! Debug render command for TUI development +//! +//! Takes JSON state as input and outputs a single rendered frame as text. +//! Useful for debugging view model derivation and rendering without running the full TUI. + +use eyre::{Context, Result}; +use ratatui::{Terminal, backend::TestBackend}; +use serde::Deserialize; +use std::io::{self, Read}; +use std::time::Instant; + +use crate::tui::{ + render::{RenderContext, render}, + state::{AppMode, AppState, ConversationEvent, StreamingStatus}, + view_model::Blocks, +}; + +/// JSON input format for debug rendering +#[derive(Debug, Deserialize)] +pub struct DebugInput { + /// Conversation events in API format + pub events: Vec, + /// Current mode: "Input", "Generating", "Streaming", "Review", "Error" + #[serde(default = "default_mode")] + pub mode: String, + /// Text being streamed (for Streaming mode) + #[serde(default)] + pub streaming_text: String, + /// Current input buffer + #[serde(default)] + pub input: String, + /// Cursor position + #[serde(default)] + pub cursor_pos: usize, + /// Spinner frame (0-3) + #[serde(default)] + pub spinner_frame: usize, + /// Error message + #[serde(default)] + pub error: Option, + /// Session ID from server + #[serde(default)] + pub session_id: Option, + /// Streaming status + #[serde(default)] + pub streaming_status: Option, + /// Whether current turn was interrupted + #[serde(default)] + pub was_interrupted: bool, + /// Terminal width for rendering + #[serde(default = "default_width")] + pub width: u16, + /// Terminal height for rendering + #[serde(default = "default_height")] + pub height: u16, +} + +fn default_mode() -> String { + "Review".to_string() +} + +fn default_width() -> u16 { + 80 +} + +fn default_height() -> u16 { + // Default to a reasonable height; state files include calculated height + 50 +} + +/// Event input matching the API protocol format +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum EventInput { + UserMessage { + content: String, + }, + Text { + content: String, + }, + ToolCall { + id: String, + name: String, + input: serde_json::Value, + }, + ToolResult { + tool_use_id: String, + content: String, + #[serde(default)] + is_error: bool, + }, +} + +impl From for ConversationEvent { + fn from(input: EventInput) -> Self { + match input { + EventInput::UserMessage { content } => ConversationEvent::UserMessage { content }, + EventInput::Text { content } => ConversationEvent::Text { content }, + EventInput::ToolCall { id, name, input } => { + ConversationEvent::ToolCall { id, name, input } + } + EventInput::ToolResult { + tool_use_id, + content, + is_error, + } => ConversationEvent::ToolResult { + tool_use_id, + content, + is_error, + }, + } + } +} + +impl DebugInput { + /// Parse JSON from string + pub fn from_json(json: &str) -> Result { + serde_json::from_str(json).context("Failed to parse debug input JSON") + } + + /// Convert to AppState + pub fn to_state(&self) -> AppState { + let mode = match self.mode.as_str() { + "Input" => AppMode::Input, + "Generating" => AppMode::Generating, + "Streaming" => AppMode::Streaming, + "Review" => AppMode::Review, + "Error" => AppMode::Error, + _ => AppMode::Review, + }; + + let events: Vec = self.events.iter().cloned().map(Into::into).collect(); + + let streaming_status = self + .streaming_status + .as_ref() + .map(|s| StreamingStatus::from_status_str(s)); + + // Create textarea from input and set cursor position + let mut textarea = tui_textarea::TextArea::from(self.input.lines()); + // Disable underline on cursor line + textarea.set_cursor_line_style(ratatui::style::Style::default()); + // Enable word wrapping + textarea.set_wrap_mode(tui_textarea::WrapMode::Word); + // Note: cursor_pos from old format is character-based; new format has row/col + // For compatibility, just move to end if we have text + if !self.input.is_empty() { + textarea.move_cursor(tui_textarea::CursorMove::End); + } + + AppState { + mode, + events, + streaming_text: self.streaming_text.clone(), + textarea, + error: self.error.clone(), + should_exit: false, + exit_action: None, + session_id: self.session_id.clone(), + streaming_status, + was_interrupted: self.was_interrupted, + spinner_frame: self.spinner_frame, + last_spinner_tick: Instant::now(), + streaming_started: None, + confirmation_pending: false, + } + } +} + +/// Output format options +#[derive(Debug, Clone, Copy, Default)] +pub enum OutputFormat { + /// Raw terminal output (ANSI) + #[default] + Ansi, + /// Plain text (strips ANSI codes) + Plain, + /// JSON with blocks structure + Json, +} + +/// Run the debug render command +pub async fn run(input_file: Option, format: OutputFormat) -> Result<()> { + // Read input JSON + let json = if let Some(path) = input_file { + std::fs::read_to_string(&path).context(format!("Failed to read input file: {}", path))? + } else { + let mut buffer = String::new(); + io::stdin() + .read_to_string(&mut buffer) + .context("Failed to read from stdin")?; + buffer + }; + + let debug_input = DebugInput::from_json(&json)?; + let state = debug_input.to_state(); + + match format { + OutputFormat::Json => { + // Output the derived blocks as JSON + let blocks = Blocks::from_state(&state); + println!( + "{}", + serde_json::to_string_pretty(&blocks_to_json(&blocks))? + ); + } + OutputFormat::Plain | OutputFormat::Ansi => { + // Render to a test backend + let backend = TestBackend::new(debug_input.width, debug_input.height); + let mut terminal = Terminal::new(backend)?; + + // Load default theme + let settings = atuin_client::settings::Settings::new()?; + let mut theme_manager = atuin_client::theme::ThemeManager::new(None, None); + let theme = theme_manager.load_theme(&settings.theme.name, None); + + let ctx = RenderContext { + theme, + anchor_col: 0, + textarea: Some(&state.textarea), + max_height: debug_input.height, + }; + + terminal.draw(|frame| { + render(frame, &state, &ctx); + })?; + + // Get buffer content + let buffer = terminal.backend().buffer(); + let output = buffer_to_string(buffer, matches!(format, OutputFormat::Plain)); + print!("{}", output); + } + } + + Ok(()) +} + +/// Convert blocks to JSON for debugging +fn blocks_to_json(blocks: &Blocks) -> serde_json::Value { + serde_json::json!({ + "count": blocks.items.len(), + "blocks": blocks.items.iter().map(|block| { + serde_json::json!({ + "separator_above": block.separator_above, + "title": block.title, + "content": block.content.iter().map(content_to_json).collect::>() + }) + }).collect::>() + }) +} + +fn content_to_json(content: &crate::tui::view_model::Content) -> serde_json::Value { + use crate::tui::view_model::Content; + match content { + Content::Input { + text, + active, + cursor_pos, + } => serde_json::json!({ + "type": "Input", + "text": text, + "active": active, + "cursor_pos": cursor_pos + }), + Content::Command { text, faded } => serde_json::json!({ + "type": "Command", + "text": text, + "faded": faded + }), + Content::Text { markdown } => serde_json::json!({ + "type": "Text", + "markdown": markdown + }), + Content::Error { message } => serde_json::json!({ + "type": "Error", + "message": message + }), + Content::Warning { + kind, + text, + pending_confirm, + } => serde_json::json!({ + "type": "Warning", + "kind": format!("{:?}", kind), + "text": text, + "pending_confirm": pending_confirm + }), + Content::Spinner { frame, status_text } => serde_json::json!({ + "type": "Spinner", + "frame": frame, + "status_text": status_text + }), + Content::ToolStatus { + completed_count, + current_label, + frame, + } => serde_json::json!({ + "type": "ToolStatus", + "completed_count": completed_count, + "current_label": current_label, + "frame": frame + }), + } +} + +/// Convert ratatui buffer to string +fn buffer_to_string(buffer: &ratatui::buffer::Buffer, strip_ansi: bool) -> String { + let area = buffer.area; + let mut output = String::new(); + + for y in 0..area.height { + for x in 0..area.width { + let cell = &buffer[(x, y)]; + if strip_ansi { + output.push_str(cell.symbol()); + } else { + // Include ANSI styling + let fg = cell.fg; + let bg = cell.bg; + let mods = cell.modifier; + + // Simple ANSI encoding + if fg != ratatui::style::Color::Reset + || bg != ratatui::style::Color::Reset + || !mods.is_empty() + { + output.push_str("\x1b["); + let mut first = true; + + if mods.contains(ratatui::style::Modifier::BOLD) { + output.push('1'); + first = false; + } + if mods.contains(ratatui::style::Modifier::DIM) { + if !first { + output.push(';'); + } + output.push('2'); + first = false; + } + if mods.contains(ratatui::style::Modifier::REVERSED) { + if !first { + output.push(';'); + } + output.push('7'); + first = false; + } + if mods.contains(ratatui::style::Modifier::UNDERLINED) { + if !first { + output.push(';'); + } + output.push('4'); + first = false; + } + + if let Some(code) = color_to_ansi(fg, true) { + if !first { + output.push(';'); + } + output.push_str(&code); + first = false; + } + + if let Some(code) = color_to_ansi(bg, false) { + if !first { + output.push(';'); + } + output.push_str(&code); + } + + output.push('m'); + } + + output.push_str(cell.symbol()); + + if fg != ratatui::style::Color::Reset + || bg != ratatui::style::Color::Reset + || !mods.is_empty() + { + output.push_str("\x1b[0m"); + } + } + } + output.push('\n'); + } + + output +} + +fn color_to_ansi(color: ratatui::style::Color, foreground: bool) -> Option { + use ratatui::style::Color; + let base = if foreground { 30 } else { 40 }; + + match color { + Color::Reset => None, + Color::Black => Some((base).to_string()), + Color::Red => Some((base + 1).to_string()), + Color::Green => Some((base + 2).to_string()), + Color::Yellow => Some((base + 3).to_string()), + Color::Blue => Some((base + 4).to_string()), + Color::Magenta => Some((base + 5).to_string()), + Color::Cyan => Some((base + 6).to_string()), + Color::Gray | Color::White => Some((base + 7).to_string()), + Color::DarkGray => Some((base + 60).to_string()), + Color::LightRed => Some((base + 61).to_string()), + Color::LightGreen => Some((base + 62).to_string()), + Color::LightYellow => Some((base + 63).to_string()), + Color::LightBlue => Some((base + 64).to_string()), + Color::LightMagenta => Some((base + 65).to_string()), + Color::LightCyan => Some((base + 66).to_string()), + Color::Indexed(i) => Some(format!("{}8;5;{}", if foreground { 3 } else { 4 }, i)), + Color::Rgb(r, g, b) => Some(format!( + "{}8;2;{};{};{}", + if foreground { 3 } else { 4 }, + r, + g, + b + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_simple_input() { + let json = r#"{ + "events": [ + {"type": "user_message", "content": "list files"}, + {"type": "tool_call", "id": "123", "name": "suggest_command", "input": {"command": "ls -la"}} + ], + "mode": "Review" + }"#; + + let input = DebugInput::from_json(json).unwrap(); + assert_eq!(input.events.len(), 2); + assert_eq!(input.mode, "Review"); + + let state = input.to_state(); + assert_eq!(state.events.len(), 2); + assert_eq!(state.mode, AppMode::Review); + } + + #[test] + fn test_parse_streaming_state() { + let json = r#"{ + "events": [ + {"type": "user_message", "content": "explain flags"} + ], + "mode": "Streaming", + "streaming_text": "The -l flag means..." + }"#; + + let input = DebugInput::from_json(json).unwrap(); + let state = input.to_state(); + assert_eq!(state.mode, AppMode::Streaming); + assert_eq!(state.streaming_text, "The -l flag means..."); + } +} diff --git a/crates/atuin-ai/src/commands/init.rs b/crates/atuin-ai/src/commands/init.rs index bf5c6256..8174b583 100644 --- a/crates/atuin-ai/src/commands/init.rs +++ b/crates/atuin-ai/src/commands/init.rs @@ -1,9 +1,29 @@ -pub async fn run() -> eyre::Result<()> { - let zsh_function = generate_zsh_integration(); - println!("{}", zsh_function); +use crate::commands::detect_shell; + +pub async fn run(shell: String) -> eyre::Result<()> { + let integration = match shell.as_str() { + "zsh" => generate_zsh_integration(), + "bash" => generate_bash_integration(), + "fish" => generate_fish_integration(), + "auto" => generate_auto_integration()?, + _ => eyre::bail!("Unsupported shell: {}", shell), + }; + + println!("{}", integration); Ok(()) } +fn generate_auto_integration() -> eyre::Result<&'static str> { + let shell = detect_shell(); + match shell.as_deref() { + Some("zsh") => Ok(generate_zsh_integration()), + Some("bash") => Ok(generate_bash_integration()), + Some("fish") => Ok(generate_fish_integration()), + Some(s) => eyre::bail!("Unsupported shell: {}", s), + None => eyre::bail!("Could not detect shell"), + } +} + /// Generate the zsh integration function - pure function for easy testing pub fn generate_zsh_integration() -> &'static str { r#" @@ -53,6 +73,111 @@ bindkey '?' _atuin_ai_question_mark # Question mark .trim() } +/// Generate the bash integration function - pure function for easy testing +pub fn generate_bash_integration() -> &'static str { + r#" +# Question mark at start of line - natural language mode +_atuin_ai_question_mark() { + # If buffer is empty or just contains '?', trigger natural language mode + if [[ -z "$READLINE_LINE" || "$READLINE_LINE" == "?" ]]; then + READLINE_LINE="" + READLINE_POINT=0 + + local output + output=$(atuin-ai inline --natural-language 3>&1 1>&2 2>&3) + + if [[ $output == __atuin_ai_cancel__ ]]; then + # User cancelled, do nothing + READLINE_LINE="" + READLINE_POINT=0 + elif [[ $output == __atuin_ai_execute__:* ]]; then + # Execute the command immediately + READLINE_LINE=${output#__atuin_ai_execute__:} + READLINE_POINT=${#READLINE_LINE} + # Note: We can't directly execute in bash bind -x, but we can + # use a workaround by binding to a macro that accepts the line + bind '"\C-x\C-a": accept-line' + bind -x '"\C-x\C-e": _atuin_ai_question_mark' + elif [[ $output == __atuin_ai_insert__:* ]]; then + # Insert the command for editing + READLINE_LINE=${output#__atuin_ai_insert__:} + READLINE_POINT=${#READLINE_LINE} + elif [[ -n $output ]]; then + # Default: insert for editing + READLINE_LINE=$output + READLINE_POINT=${#READLINE_LINE} + fi + else + # Not at empty prompt, just insert the question mark + READLINE_LINE="${READLINE_LINE:0:READLINE_POINT}?${READLINE_LINE:READLINE_POINT}" + ((READLINE_POINT++)) + fi +} + +# Set up keybindings +# Bash requires special handling: we use bind -x for the function, +# but need a two-step approach for execute mode +__atuin_ai_accept_line="" + +_atuin_ai_question_mark_wrapper() { + _atuin_ai_question_mark + if [[ -n "$__atuin_ai_accept_line" ]]; then + __atuin_ai_accept_line="" + fi +} + +bind -x '"?": _atuin_ai_question_mark' +"# + .trim() +} + +/// Generate the fish integration function - pure function for easy testing +pub fn generate_fish_integration() -> &'static str { + r#" +# Question mark at start of line - natural language mode +function _atuin_ai_question_mark + set -l buf (commandline -b) + + # If buffer is empty or just contains '?', trigger natural language mode + if test -z "$buf" -o "$buf" = "?" + commandline -r "" + + # Run atuin-ai inline, swapping stdout and stderr + set -l output (atuin-ai inline --natural-language 3>&1 1>&2 2>&3 | string collect) + + if test "$output" = "__atuin_ai_cancel__" + # User cancelled, do nothing + commandline -f repaint + else if string match --quiet '__atuin_ai_execute__:*' "$output" + # Execute the command immediately + set -l cmd (string replace "__atuin_ai_execute__:" "" -- "$output" | string collect) + commandline -r "$cmd" + commandline -f repaint + commandline -f execute + else if string match --quiet '__atuin_ai_insert__:*' "$output" + # Insert the command for editing + set -l cmd (string replace "__atuin_ai_insert__:" "" -- "$output" | string collect) + commandline -r "$cmd" + commandline -f repaint + else if test -n "$output" + # Default: insert for editing + commandline -r "$output" + commandline -f repaint + else + commandline -f repaint + end + else + # Not at empty prompt, just insert the question mark + commandline -i "?" + end +end + +# Set up keybindings +bind "?" _atuin_ai_question_mark +"# + .trim() +} + #[cfg(test)] mod tests { use super::*; @@ -67,4 +192,28 @@ mod tests { assert!(result.contains("__atuin_ai_execute__")); assert!(result.contains("__atuin_ai_insert__")); } + + #[test] + fn test_generate_bash_integration() { + let result = generate_bash_integration(); + assert!(result.contains("_atuin_ai_question_mark")); + assert!(result.contains("bind")); + assert!(result.contains("READLINE_LINE")); + assert!(result.contains("atuin-ai inline")); + assert!(result.contains("__atuin_ai_cancel__")); + assert!(result.contains("__atuin_ai_execute__")); + assert!(result.contains("__atuin_ai_insert__")); + } + + #[test] + fn test_generate_fish_integration() { + let result = generate_fish_integration(); + assert!(result.contains("_atuin_ai_question_mark")); + assert!(result.contains("bind")); + assert!(result.contains("commandline")); + assert!(result.contains("atuin-ai inline")); + assert!(result.contains("__atuin_ai_cancel__")); + assert!(result.contains("__atuin_ai_execute__")); + assert!(result.contains("__atuin_ai_insert__")); + } } diff --git a/crates/atuin-ai/src/commands/inline.rs b/crates/atuin-ai/src/commands/inline.rs index cfa27db4..3f9278a2 100644 --- a/crates/atuin-ai/src/commands/inline.rs +++ b/crates/atuin-ai/src/commands/inline.rs @@ -1,52 +1,52 @@ +use crate::commands::detect_shell; +use crate::tui::render::render; +use crate::tui::{ + App, AppEvent, AppMode, ConversationEvent, EventLoop, ExitAction, RenderContext, TerminalGuard, + calculate_needed_height, install_panic_hook, +}; +use atuin_client::theme::ThemeManager; use atuin_common::tls::ensure_crypto_provider; use crossterm::{ - cursor, event::{self, Event, KeyCode}, terminal::{disable_raw_mode, enable_raw_mode}, }; +use eventsource_stream::Eventsource; use eyre::{Context as _, Result, bail}; -use ratatui::{ - Frame, Terminal, TerminalOptions, Viewport, - backend::CrosstermBackend, - layout::{Alignment, Rect}, - text::Line, - widgets::{Block, Borders, Paragraph, Wrap}, -}; +use futures::StreamExt; use reqwest::Url; -use serde::{Deserialize, Serialize}; -use std::time::Duration; - -#[derive(Debug, Serialize)] -struct GenerateRequest { - query: String, - description: String, - context: GenerateContext, -} - -#[derive(Debug, Serialize)] -struct GenerateContext { - os: String, - shell: String, - pwd: Option, -} - -#[derive(Debug, Deserialize)] -struct GenerateResponse { - command: String, - #[serde(default)] - explanation: Option, -} +use std::io::Write; pub async fn run( initial_command: Option, natural_language: bool, api_endpoint: Option, + api_token: Option, + keep_output: bool, + debug_state_file: Option, ) -> Result<()> { + // Install panic hook once at entry point to ensure terminal restoration + install_panic_hook(); + + // Token and endpoint priority: + // 1. Command line arguments/environment variables + // 2. Settings file + // 3. Default let settings = atuin_client::settings::Settings::new()?; - let endpoint = api_endpoint - .as_deref() - .unwrap_or(settings.hub_address.as_str()); - let token = ensure_hub_session(&settings, endpoint).await?; + let endpoint = api_endpoint.as_deref().unwrap_or( + settings + .ai + .ai_endpoint + .as_deref() + .unwrap_or("https://hub.atuin.sh"), + ); + let api_token = api_token.as_deref().or(settings.ai.ai_api_token.as_deref()); + + let token = if let Some(token) = &api_token { + token.to_string() + } else { + ensure_hub_session(&settings, endpoint).await? + }; + let action = run_inline_tui( endpoint.to_string(), token, @@ -55,6 +55,8 @@ pub async fn run( } else { initial_command }, + keep_output, + debug_state_file, ) .await?; emit_shell_result(action.0, &action.1); @@ -95,55 +97,172 @@ async fn ensure_hub_session( Ok(token) } -async fn generate_command( - hub_address: &str, - token: &str, - description: &str, -) -> Result { - ensure_crypto_provider(); - let endpoint = hub_url(hub_address, "/api/cli/generate")?; - let request = GenerateRequest { - query: description.to_string(), - description: description.to_string(), - context: GenerateContext { - os: detect_os(), - shell: detect_shell(), - pwd: std::env::current_dir() - .ok() - .map(|path| path.to_string_lossy().into_owned()), - }, - }; +/// SSE event received from chat endpoint +#[derive(Debug, Clone)] +enum ChatStreamEvent { + /// Text chunk to display + TextChunk(String), + /// Tool call event (need to echo back, may contain suggest_command) + ToolCall { + id: String, + name: String, + input: serde_json::Value, + }, + /// Tool result from server-side execution + ToolResult { + tool_use_id: String, + content: String, + is_error: bool, + }, + /// Status update from server + Status(String), + /// Stream complete + Done { session_id: String }, + /// Error from server + Error(String), +} + +fn create_chat_stream( + hub_address: String, + token: String, + session_id: Option, + messages: Vec, + settings: &atuin_client::settings::Settings, +) -> std::pin::Pin> + Send>> { + let send_cwd = settings.ai.send_cwd; + + Box::pin(async_stream::stream! { + ensure_crypto_provider(); + let endpoint = match hub_url(&hub_address, "/api/cli/chat") { + Ok(url) => url, + Err(e) => { + yield Err(e); + return; + } + }; + + // Build request body + let mut request_body = serde_json::json!({ + "messages": messages, + "context": { + "os": detect_os(), + "shell": detect_shell(), + "pwd": if send_cwd { std::env::current_dir() + .ok() + .map(|path| path.to_string_lossy().into_owned()) } else { None }, + } + }); + + // Include session_id only if present (not on first request) + if let Some(ref sid) = session_id { + request_body["session_id"] = serde_json::json!(sid); + } - let client = reqwest::Client::new(); - let response = client - .post(endpoint) - .bearer_auth(token) - .json(&request) - .send() - .await - .context("failed to call Atuin Hub generate endpoint")?; - - if response.status().is_success() { - let generated = response - .json::() + + let client = reqwest::Client::new(); + let response = match client + .post(endpoint.clone()) + .header("Accept", "text/event-stream") + .bearer_auth(&token) + .json(&request_body) + .send() .await - .context("failed to decode generate response")?; + { + Ok(resp) => resp, + Err(e) => { + yield Err(eyre::eyre!("Failed to send SSE request: {}", e)); + return; + } + }; - if generated.command.trim().is_empty() { - bail!("Hub returned an empty command. Please try again with a more specific request."); + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED { + // Clear saved session on auth error + let _ = atuin_client::hub::delete_session().await; + yield Err(eyre::eyre!("Hub session expired. Re-run to authenticate again.")); + return; + } + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + yield Err(eyre::eyre!("SSE request failed ({}): {}", status, body)); + return; } - return Ok(generated); - } + let byte_stream = response.bytes_stream(); + let mut stream = byte_stream.eventsource(); - if response.status() == reqwest::StatusCode::UNAUTHORIZED { - atuin_client::hub::delete_session().await?; - bail!("Hub session expired. Re-run to authenticate again."); - } + while let Some(event) = stream.next().await { + match event { + Ok(sse_event) => { + let event_type = sse_event.event.as_str(); + let data = sse_event.data.clone(); + + tracing::debug!(event_type = %event_type, data = %data, "SSE event received"); - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - bail!("Hub request failed ({status}): {body}"); + match event_type { + "text" => { + if let Ok(json) = serde_json::from_str::(&data) + && let Some(content) = json.get("content").and_then(|v| v.as_str()) + { + yield Ok(ChatStreamEvent::TextChunk(content.to_string())); + } + } + "tool_call" => { + if let Ok(json) = serde_json::from_str::(&data) { + let id = json.get("id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let name = json.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let input = json.get("input").cloned().unwrap_or(serde_json::json!({})); + yield Ok(ChatStreamEvent::ToolCall { id, name, input }); + } + } + "tool_result" => { + if let Ok(json) = serde_json::from_str::(&data) { + let tool_use_id = json.get("tool_use_id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let content = json.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let is_error = json.get("is_error").and_then(|v| v.as_bool()).unwrap_or(false); + yield Ok(ChatStreamEvent::ToolResult { tool_use_id, content, is_error }); + } + } + "status" => { + if let Ok(json) = serde_json::from_str::(&data) + && let Some(state) = json.get("state").and_then(|v| v.as_str()) + { + yield Ok(ChatStreamEvent::Status(state.to_string())); + } + } + "done" => { + if let Ok(json) = serde_json::from_str::(&data) { + let session_id = json.get("session_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + yield Ok(ChatStreamEvent::Done { session_id }); + } else { + yield Ok(ChatStreamEvent::Done { session_id: String::new() }); + } + break; + } + "error" => { + if let Ok(json) = serde_json::from_str::(&data) { + let message = json.get("message").and_then(|v| v.as_str()).unwrap_or("Unknown error").to_string(); + yield Ok(ChatStreamEvent::Error(message)); + } else { + yield Ok(ChatStreamEvent::Error(data)); + } + break; + } + _ => { + // Unknown event type, ignore + } + } + } + Err(e) => { + yield Err(eyre::eyre!("SSE error: {}", e)); + break; + } + } + } + }) } fn hub_url(base: &str, path: &str) -> Result { @@ -162,35 +281,11 @@ fn detect_os() -> String { match std::env::consts::OS { "macos" => "macos".to_string(), "linux" => "linux".to_string(), + "windows" => "windows".to_string(), _ => "linux".to_string(), } } -fn detect_shell() -> String { - if let Ok(shell) = std::env::var("ATUIN_SHELL") - && !shell.trim().is_empty() - { - return shell; - } - - let shell = std::env::var("SHELL") - .ok() - .and_then(|value| { - std::path::Path::new(&value) - .file_name() - .map(std::ffi::OsStr::to_string_lossy) - .map(std::borrow::Cow::into_owned) - }) - .filter(|value| !value.trim().is_empty()); - - match shell.as_deref() { - Some("zsh") => "zsh".to_string(), - Some("fish") => "fish".to_string(), - Some("bash") => "bash".to_string(), - _ => "bash".to_string(), - } -} - #[derive(Clone, Copy)] enum Action { Execute, @@ -198,105 +293,306 @@ enum Action { Cancel, } +/// Serialize AppState to JSON for debug logging +fn state_to_json(state: &crate::tui::AppState) -> serde_json::Value { + let events: Vec = state.events.iter().map(|e| e.to_json()).collect(); + + let mode = match state.mode { + AppMode::Input => "Input", + AppMode::Generating => "Generating", + AppMode::Streaming => "Streaming", + AppMode::Review => "Review", + AppMode::Error => "Error", + }; + + // Get input and cursor from textarea + let input = state.input(); + let cursor = state.textarea.cursor(); + + let mut json = serde_json::json!({ + "events": events, + "mode": mode, + "input": input, + "cursor_row": cursor.0, + "cursor_col": cursor.1, + "spinner_frame": state.spinner_frame, + "confirmation_pending": state.confirmation_pending, + }); + + // Add streaming fields if in streaming mode + if !state.streaming_text.is_empty() { + json["streaming_text"] = serde_json::json!(state.streaming_text); + } + if let Some(ref status) = state.streaming_status { + json["streaming_status"] = serde_json::json!(status.display_text()); + } + if let Some(ref err) = state.error { + json["error"] = serde_json::json!(err); + } + + json +} + +/// Debug logger that writes state changes to a file +struct DebugStateLogger { + file: std::fs::File, + entry_count: usize, + width: u16, +} + +impl DebugStateLogger { + fn new(path: &str) -> Result { + let file = std::fs::File::create(path) + .with_context(|| format!("Failed to create debug state file: {}", path))?; + // Get terminal width, default to 80 + let (width, _) = crossterm::terminal::size().unwrap_or((80, 24)); + Ok(Self { + file, + entry_count: 0, + width, + }) + } + + fn log(&mut self, label: &str, state: &crate::tui::AppState) { + use crate::tui::calculate_needed_height; + + self.entry_count += 1; + let timestamp_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis()) + .unwrap_or(0); + + // Calculate the actual content height needed for this state + let content_height = calculate_needed_height(state); + + let mut state_json = state_to_json(state); + // Add dimensions for accurate replay + state_json["width"] = serde_json::json!(self.width); + state_json["height"] = serde_json::json!(content_height); + + let entry = serde_json::json!({ + "entry": self.entry_count, + "label": label, + "timestamp_ms": timestamp_ms, + "state": state_json, + }); + + // Write as JSONL (one JSON object per line) + if let Err(e) = writeln!(self.file, "{}", entry) { + tracing::warn!("Failed to write debug state: {}", e); + } + let _ = self.file.flush(); + } +} + async fn run_inline_tui( endpoint: String, token: String, initial_prompt: Option, + keep_output: bool, + debug_state_file: Option, ) -> Result<(Action, String)> { - let mut ui = InlineUi::new()?; - let mut prompt = initial_prompt.unwrap_or_default(); - let mut spinner_idx = 0usize; + // Initialize terminal guard and app state + let mut guard = TerminalGuard::new(keep_output)?; + let mut app = App::new(); + if let Some(prompt) = initial_prompt { + // Set initial text in textarea + let mut textarea = tui_textarea::TextArea::from(prompt.lines()); + // Disable underline on cursor line + textarea.set_cursor_line_style(ratatui::style::Style::default()); + // Enable word wrapping + textarea.set_wrap_mode(tui_textarea::WrapMode::Word); + // Move cursor to end + textarea.move_cursor(tui_textarea::CursorMove::End); + app.state.textarea = textarea; + } - loop { - ui.render_prompt(&prompt)?; - if !event::poll(Duration::from_millis(250)).context("failed to poll for input")? { - continue; - } + // Initialize debug state logger if requested + let mut debug_logger = debug_state_file + .map(|path| DebugStateLogger::new(&path)) + .transpose()?; - let ev = event::read().context("failed to read terminal event")?; - let Event::Key(key) = ev else { - continue; + // Helper macro to log state changes + macro_rules! log_state { + ($label:expr) => { + if let Some(ref mut logger) = debug_logger { + logger.log($label, &app.state); + } }; + } - match key.code { - KeyCode::Esc => return Ok((Action::Cancel, String::new())), - KeyCode::Backspace => { - prompt.pop(); - } - KeyCode::Enter => { - let query = prompt.trim().to_string(); - if query.is_empty() { - return Ok((Action::Cancel, String::new())); - } + // Log initial state + log_state!("init"); - let response = loop { - let endpoint_clone = endpoint.clone(); - let token_clone = token.clone(); - let query_clone = query.clone(); - let task = tokio::spawn(async move { - generate_command(&endpoint_clone, &token_clone, &query_clone).await - }); - - let generated = loop { - if task.is_finished() { - break task.await.context("generate task join failed")?; - } + // Load theme + let settings = atuin_client::settings::Settings::new()?; + let mut theme_manager = ThemeManager::new(None, None); + let theme = theme_manager.load_theme(&settings.theme.name, None); - ui.render_generating(&prompt, spinner_idx)?; - spinner_idx = (spinner_idx + 1) % SPINNER_FRAMES.len(); + // Initialize event loop + let mut event_loop = EventLoop::new(); - if event::poll(Duration::from_millis(100)) - .context("failed to poll while generating")? - { - let ev = event::read().context("failed reading generate event")?; - if let Event::Key(key) = ev - && key.code == KeyCode::Esc - { - task.abort(); - return Ok((Action::Cancel, String::new())); + // Track chat stream + let mut chat_stream: Option< + std::pin::Pin> + Send>>, + > = None; + + loop { + // Ensure viewport is large enough for current content (capped at terminal height) + let needed_height = calculate_needed_height(&app.state); + let actual_height = guard.ensure_height(needed_height)?; + + // Render current state + let anchor_col = guard.anchor_col(); + let ctx = RenderContext { + theme, + anchor_col, + textarea: Some(&app.state.textarea), + max_height: actual_height, + }; + // Handle draw errors gracefully - cursor position reads can fail during resize + if let Err(e) = guard.terminal().draw(|frame| { + render(frame, &app.state, &ctx); + }) { + let err_msg = e.to_string(); + if err_msg.contains("cursor position") { + // Cursor position read failed (common during terminal resize) + // Skip this frame and continue - next frame will likely succeed + tracing::debug!( + "Skipping frame due to cursor position read error: {}", + err_msg + ); + continue; + } + return Err(e.into()); + } + + // Get next event + let event = event_loop.run().await?; + + // Handle event based on app mode + match event { + AppEvent::Key(key) => { + app.handle_key(key); + log_state!("key"); + } + AppEvent::Tick => { + app.state.tick(); + + // Poll chat stream if active - keep polling until done regardless of mode + // (mode may change to Review before we receive the done event with session_id) + if let Some(stream) = &mut chat_stream { + let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref()); + match stream.as_mut().poll_next(&mut cx) { + std::task::Poll::Ready(Some(Ok(event))) => match event { + ChatStreamEvent::TextChunk(text) => { + tracing::debug!(text = %text, "Processing TextChunk"); + app.state.append_streaming_text(&text); + log_state!("text_chunk"); } - } - }; - - match generated { - Ok(value) => break value, - Err(err) => { - ui.render_error(&prompt, &err.to_string())?; - if !wait_for_retry_or_cancel()? { - return Ok((Action::Cancel, String::new())); + ChatStreamEvent::ToolCall { id, name, input } => { + tracing::debug!(id = %id, name = %name, "Processing ToolCall"); + app.state.add_tool_call(id, name, input); + log_state!("tool_call"); + } + ChatStreamEvent::ToolResult { + tool_use_id, + content, + is_error, + } => { + tracing::debug!(tool_use_id = %tool_use_id, "Processing ToolResult"); + app.state.add_tool_result(tool_use_id, content, is_error); + log_state!("tool_result"); + } + ChatStreamEvent::Status(status) => { + tracing::debug!(status = %status, "Processing Status"); + app.state.update_streaming_status(&status); + log_state!("status"); + } + ChatStreamEvent::Done { session_id } => { + tracing::debug!(session_id = %session_id, "Processing Done"); + chat_stream = None; + if !session_id.is_empty() { + app.state.store_session_id(session_id); + } + app.state.finalize_streaming(); + log_state!("done"); } + ChatStreamEvent::Error(msg) => { + tracing::debug!(error = %msg, "Processing Error"); + chat_stream = None; + app.state.streaming_error(msg); + log_state!("error"); + } + }, + std::task::Poll::Ready(Some(Err(e))) => { + chat_stream = None; + app.state.streaming_error(e.to_string()); + log_state!("stream_error"); } - } - }; - - loop { - ui.render_review(&prompt, &response)?; - if !event::poll(Duration::from_millis(250)) - .context("failed to poll in review")? - { - continue; - } - - let ev = event::read().context("failed to read review event")?; - let Event::Key(key) = ev else { - continue; - }; - - match key.code { - KeyCode::Enter => return Ok((Action::Execute, response.command)), - KeyCode::Tab => return Ok((Action::Insert, response.command)), - KeyCode::Esc => return Ok((Action::Cancel, String::new())), - KeyCode::Char('e') => break, - _ => {} + std::task::Poll::Ready(None) => { + chat_stream = None; + app.state.finalize_streaming(); + log_state!("stream_end"); + } + std::task::Poll::Pending => {} } } } - KeyCode::Char(c) => { - prompt.push(c); - } _ => {} } + + // Handle user cancellation (Esc during streaming) - drop the stream + if app.state.was_interrupted && chat_stream.is_some() { + tracing::debug!("User cancelled streaming, dropping chat stream"); + chat_stream = None; + app.state.was_interrupted = false; // Reset the flag + } + + // Check exit condition + if app.state.should_exit { + break; + } + + // Handle generation trigger - unified path for all turns + if app.state.mode == AppMode::Generating && chat_stream.is_none() { + // Get the last user message from events + let last_user_content = app.state.events.iter().rev().find_map(|e| { + if let ConversationEvent::UserMessage { content } = e { + Some(content.clone()) + } else { + None + } + }); + + if last_user_content.is_some() { + // Build messages in Claude API format + let messages = app.state.events_to_messages(); + + // Transition to streaming mode + app.state.start_streaming(); + log_state!("start_streaming"); + + // Start the chat stream + chat_stream = Some(create_chat_stream( + endpoint.clone(), + token.clone(), + app.state.session_id.clone(), + messages, + &settings, + )); + } + } } + + // Map exit action to return value + let result = match app.state.exit_action { + Some(ExitAction::Execute(cmd)) => (Action::Execute, cmd), + Some(ExitAction::Insert(cmd)) => (Action::Insert, cmd), + _ => (Action::Cancel, String::new()), + }; + + Ok(result) } struct RawModeGuard; @@ -330,279 +626,3 @@ fn wait_for_login_confirmation() -> Result { } } } - -fn wait_for_retry_or_cancel() -> Result { - loop { - let ev = event::read().context("failed to read retry/cancel key")?; - if let Event::Key(key) = ev { - match key.code { - KeyCode::Enter | KeyCode::Char('r') => return Ok(true), - KeyCode::Esc => return Ok(false), - _ => {} - } - } - } -} - -const SPINNER_FRAMES: [&str; 4] = ["/", "-", "\\", "|"]; - -struct InlineUi { - terminal: Terminal>, - anchor_col: u16, -} - -impl InlineUi { - fn new() -> Result { - let anchor_col = cursor::position().map(|(x, _)| x).unwrap_or(0); - enable_raw_mode().context("failed to enable raw mode for inline UI")?; - let backend = CrosstermBackend::new(std::io::stdout()); - let terminal = Terminal::with_options( - backend, - TerminalOptions { - viewport: Viewport::Inline(16), - }, - ) - .context("failed to initialize inline UI")?; - Ok(Self { - terminal, - anchor_col, - }) - } - - fn render_prompt(&mut self, prompt: &str) -> Result<()> { - self.render(Screen::Prompt { - prompt, - footer: "[Enter]: Accept [Esc]: Cancel", - }) - } - - fn render_generating(&mut self, prompt: &str, spinner_idx: usize) -> Result<()> { - self.render(Screen::Generating { - prompt, - footer: "[Esc]: Cancel", - spinner_idx, - }) - } - - fn render_review(&mut self, prompt: &str, response: &GenerateResponse) -> Result<()> { - self.render(Screen::Review { - prompt, - response, - footer: "[Enter]: Run [Tab]: Insert [e]: Edit [Esc]: Cancel", - }) - } - - fn render_error(&mut self, prompt: &str, err: &str) -> Result<()> { - self.render(Screen::Error { - prompt, - err, - footer: "[Enter]/[r]: Retry [Esc]: Cancel", - }) - } - - fn render(&mut self, screen: Screen<'_>) -> Result<()> { - self.terminal - .draw(|f| draw_screen(f, screen, self.anchor_col)) - .context("failed rendering inline UI")?; - Ok(()) - } -} - -impl Drop for InlineUi { - fn drop(&mut self) { - let _ = self.terminal.clear(); - let _ = disable_raw_mode(); - } -} - -enum Screen<'a> { - Prompt { - prompt: &'a str, - footer: &'a str, - }, - Generating { - prompt: &'a str, - footer: &'a str, - spinner_idx: usize, - }, - Review { - prompt: &'a str, - response: &'a GenerateResponse, - footer: &'a str, - }, - Error { - prompt: &'a str, - err: &'a str, - footer: &'a str, - }, -} - -fn draw_screen(frame: &mut Frame, screen: Screen<'_>, anchor_col: u16) { - let area = frame.area(); - let desired_width = 64u16.min(area.width.saturating_sub(2)).max(32); - let content_width = usize::from(desired_width.saturating_sub(2)).max(1); - let (content_preview, _, _) = build_screen_content(&screen, content_width); - let desired_height = (wrapped_line_count(&content_preview, content_width) as u16) - .saturating_add(2) - .min(area.height.max(1)) - .max(3); - - let max_x = area.x + area.width.saturating_sub(desired_width); - let preferred_x = area.x + anchor_col.saturating_sub(2); - let card = Rect { - x: preferred_x.min(max_x), - y: area.y, - width: desired_width, - height: desired_height, - }; - - let footer = match &screen { - Screen::Prompt { footer, .. } - | Screen::Generating { footer, .. } - | Screen::Review { footer, .. } - | Screen::Error { footer, .. } => *footer, - }; - - let block = Block::default() - .borders(Borders::ALL) - .title("Describe the command you'd like to generate:") - .title_bottom(Line::from(footer).alignment(Alignment::Right)); - - let content_area = block.inner(card); - frame.render_widget(block, card); - - let (content, show_cursor, cursor_prompt) = - build_screen_content(&screen, usize::from(content_area.width).max(1)); - - let paragraph = Paragraph::new(content).wrap(Wrap { trim: false }); - frame.render_widget(paragraph, content_area); - - if show_cursor { - let width = usize::from(content_area.width).max(1); - let (cursor_row, cursor_col) = - prompt_cursor_position(cursor_prompt.as_deref().unwrap_or_default(), width); - let cursor_x = content_area.x.saturating_add(cursor_col); - let cursor_y = content_area.y.saturating_add(cursor_row); - frame.set_cursor_position((cursor_x, cursor_y)); - } -} - -fn format_prompt(prompt: &str) -> String { - if prompt.is_empty() { - return "> ".to_string(); - } - format!("> {prompt}") -} - -fn wrapped_line_count(text: &str, width: usize) -> usize { - if width == 0 { - return 1; - } - - text.split('\n') - .map(|line| { - let len = line.chars().count(); - len.max(1).div_ceil(width) - }) - .sum::() - .max(1) -} - -fn build_screen_content( - screen: &Screen<'_>, - content_width: usize, -) -> (String, bool, Option) { - match screen { - Screen::Prompt { prompt, .. } => { - let formatted = format_prompt(prompt); - (formatted, true, Some((*prompt).to_string())) - } - Screen::Generating { - prompt, - spinner_idx, - .. - } => ( - format!( - "{}\n\n{} Generating...", - format_prompt(prompt), - SPINNER_FRAMES[*spinner_idx] - ), - false, - None, - ), - Screen::Review { - prompt, response, .. - } => { - let separator = "─".repeat(content_width.max(1)); - let mut text = format!( - "{}\n\n{}\n\n$ {}\n", - format_prompt(prompt), - separator, - response.command - ); - if let Some(explanation) = &response.explanation { - text.push('\n'); - text.push_str(explanation); - } - (text, false, None) - } - Screen::Error { prompt, err, .. } => ( - format!("{}\n\nRequest failed:\n{}", format_prompt(prompt), err), - false, - None, - ), - } -} - -fn prompt_cursor_position(prompt: &str, width: usize) -> (u16, u16) { - if width == 0 { - return (0, 0); - } - - // The visible prompt line is always `> {prompt}`. - // We mimic word-wrapping so cursor tracking matches visual layout. - let mut row = 0usize; - let mut col = 2usize; // "> " - - let mut saw_any_word = false; - for word in prompt.split_whitespace() { - let word_len = word.chars().count(); - if !saw_any_word { - saw_any_word = true; - if col + word_len <= width { - col += word_len; - } else if word_len >= width { - let used = width.saturating_sub(col); - let remaining = word_len.saturating_sub(used); - row += 1 + (remaining / width); - col = remaining % width; - } else { - row += 1; - col = word_len; - } - continue; - } - - if col + 1 + word_len <= width { - col += 1 + word_len; - } else if word_len >= width { - row += 1 + (word_len / width); - col = word_len % width; - } else { - row += 1; - col = word_len; - } - } - - // Keep trailing spaces user typed. - let trailing_spaces = prompt.chars().rev().take_while(|c| *c == ' ').count(); - for _ in 0..trailing_spaces { - if col >= width { - row += 1; - col = 0; - } - col += 1; - } - - (row as u16, col as u16) -} diff --git a/crates/atuin-ai/src/main.rs b/crates/atuin-ai/src/main.rs index 6302bbda..fb1e517e 100644 --- a/crates/atuin-ai/src/main.rs +++ b/crates/atuin-ai/src/main.rs @@ -1,4 +1,5 @@ pub mod commands; +pub mod tui; #[tokio::main] async fn main() -> eyre::Result<()> { diff --git a/crates/atuin-ai/src/tui/app.rs b/crates/atuin-ai/src/tui/app.rs new file mode 100644 index 00000000..ecb1eb81 --- /dev/null +++ b/crates/atuin-ai/src/tui/app.rs @@ -0,0 +1,157 @@ +use super::state::{AppMode, AppState, ExitAction}; +use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; +use tui_textarea::{Input, Key}; + +/// Thin wrapper around AppState for compatibility +/// All state lives in AppState, this just provides the handle_key interface +pub struct App { + pub state: AppState, +} + +impl App { + pub fn new() -> Self { + Self { + state: AppState::new(), + } + } + + /// Handle a key event. Returns true if render is needed. + pub fn handle_key(&mut self, key: KeyEvent) -> bool { + match self.state.mode { + AppMode::Input => self.handle_input_key(key), + AppMode::Generating => self.handle_generating_key(key), + AppMode::Streaming => self.handle_streaming_key(key), + AppMode::Review => self.handle_review_key(key), + AppMode::Error => self.handle_error_key(key), + } + } + + fn handle_input_key(&mut self, key: KeyEvent) -> bool { + // Handle special keys ourselves + match key.code { + KeyCode::Esc => { + self.state.exit(ExitAction::Cancel); + return true; + } + KeyCode::Enter => { + if self.state.input_is_empty() { + self.state.exit(ExitAction::Cancel); + } else { + self.state.start_generating(); + } + return true; + } + _ => {} + } + + // Delegate all other keys to textarea + // Manually convert crossterm KeyEvent to tui-textarea Input + // (needed due to crossterm version mismatch) + let tui_key = match key.code { + KeyCode::Char(c) => Key::Char(c), + KeyCode::Backspace => Key::Backspace, + KeyCode::Delete => Key::Delete, + KeyCode::Left => Key::Left, + KeyCode::Right => Key::Right, + KeyCode::Up => Key::Up, + KeyCode::Down => Key::Down, + KeyCode::Home => Key::Home, + KeyCode::End => Key::End, + KeyCode::PageUp => Key::PageUp, + KeyCode::PageDown => Key::PageDown, + KeyCode::Tab => Key::Tab, + _ => Key::Null, + }; + + if tui_key != Key::Null { + let input = Input { + key: tui_key, + ctrl: key.modifiers.contains(KeyModifiers::CONTROL), + alt: key.modifiers.contains(KeyModifiers::ALT), + shift: key.modifiers.contains(KeyModifiers::SHIFT), + }; + self.state.textarea.input(input); + } + true + } + + fn handle_generating_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.cancel_generation(); + true + } + _ => false, // Discard other keys during generation + } + } + + fn handle_streaming_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.cancel_streaming(); + true + } + _ => false, // Ignore other keys during streaming + } + } + + fn handle_review_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.confirmation_pending = false; // Clear confirmation state + self.state.exit(ExitAction::Cancel); + true + } + KeyCode::Enter => { + let cmd = self.state.current_command().map(|c| c.to_string()); + if let Some(cmd) = cmd { + if self.state.is_current_command_dangerous() && !self.state.confirmation_pending + { + // First Enter on dangerous command: enter confirmation mode + self.state.confirmation_pending = true; + } else { + // Second Enter (confirmation), or non-dangerous command: execute + self.state.confirmation_pending = false; + self.state.exit(ExitAction::Execute(cmd)); + } + } + true + } + KeyCode::Tab => { + let cmd = self.state.current_command().map(|c| c.to_string()); + if let Some(cmd) = cmd { + self.state.confirmation_pending = false; // Clear on Tab too + self.state.exit(ExitAction::Insert(cmd)); + } + true + } + KeyCode::Char('f') => { + // Changed from 'e' to 'f' for follow-up mode + self.state.confirmation_pending = false; // Clear on follow-up + self.state.start_edit_mode(); + true + } + _ => false, + } + } + + fn handle_error_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.exit(ExitAction::Cancel); + true + } + KeyCode::Enter | KeyCode::Char('r') => { + self.state.retry(); + true + } + _ => false, + } + } +} + +impl Default for App { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/atuin-ai/src/tui/event.rs b/crates/atuin-ai/src/tui/event.rs new file mode 100644 index 00000000..8efbf522 --- /dev/null +++ b/crates/atuin-ai/src/tui/event.rs @@ -0,0 +1,303 @@ +use crate::tui::App; +use crossterm::event::{Event, EventStream, KeyEvent, KeyEventKind}; +use eyre::{Result, eyre}; +use futures::StreamExt; +use std::time::Duration; +use tokio::time; + +/// Base tick interval for the event loop (fast for responsive streaming) +const BASE_TICK_INTERVAL: Duration = Duration::from_millis(50); + +/// Application events that drive the TUI state machine. +/// +/// # Event Types +/// - `Key`: Keyboard input (filtered to KeyEventKind::Press only) +/// - `Tick`: Periodic event for updates (50ms base interval) +/// - `Resize`: Terminal window resize +/// - `StreamChunk/StreamDone/StreamError`: Placeholders for Phase 3 streaming +/// +/// # Design Decisions +/// - Fast 50ms base tick for responsive streaming; spinner timing handled in AppState +/// - Stream events are placeholders - will be wired to channels in Phase 3 +/// - Resize handling enables responsive layout adjustments +#[derive(Debug, Clone)] +pub enum AppEvent { + /// Keyboard input event (filtered to Press events only) + Key(KeyEvent), + + /// Periodic tick for updates (50ms base interval; spinner timing in AppState) + Tick, + + /// Terminal resize event (width, height) + Resize(u16, u16), + + /// Stream chunk received (Phase 3 placeholder) + StreamChunk(String), + + /// Stream completed successfully (Phase 3 placeholder) + StreamDone, + + /// Stream error occurred (Phase 3 placeholder) + StreamError(String), +} + +/// Async event loop that drives the TUI with prioritized event handling. +/// +/// # Priority Model (Biased Select) +/// 1. **Stream data** - Highest priority (future Phase 3 streaming) +/// 2. **Keyboard input** - Medium priority (user responsiveness) +/// 3. **Tick events** - Lowest priority (spinner animation) +/// +/// This ensures stream data is processed immediately when available, +/// keyboard input is responsive, and spinner updates don't block higher priority events. +/// +/// # Graceful Shutdown +/// - SIGINT (Ctrl+C) sets shutdown flag and breaks the loop +/// - EventStream close (stdin EOF) triggers shutdown +/// - Shutdown flag can be checked/set externally for controlled termination +/// +/// # Example +/// ```no_run +/// use atuin_ai::tui::EventLoop; +/// +/// # async fn example() -> eyre::Result<()> { +/// let mut event_loop = EventLoop::new(); +/// loop { +/// let event = event_loop.run().await?; +/// // Handle event... +/// # break; +/// } +/// # Ok(()) +/// # } +/// ``` +pub struct EventLoop { + /// Tick interval timer (created lazily on first run) + tick_timer: Option, + + /// Flag indicating a render was requested (future use in Phase 2) + #[allow(dead_code)] + render_requested: bool, + + /// Shutdown flag - when true, event loop will terminate + shutdown: bool, +} + +impl EventLoop { + /// Create a new EventLoop with default settings. + /// + /// # Defaults + /// - Tick interval: 50ms base rate (spinner timing handled separately in AppState) + /// - Render requested: false + /// - Shutdown: false + pub fn new() -> Self { + Self { + tick_timer: None, + render_requested: false, + shutdown: false, + } + } + + /// Run the event loop, returning the next application event. + /// + /// # Priority Model + /// Uses `tokio::select!` with `biased;` mode to enforce priority: + /// 1. Stream data (placeholder for Phase 3) + /// 2. Keyboard input with rapid keypress batching + /// 3. Tick for spinner animation + /// + /// # Keyboard Handling + /// - Filters to KeyEventKind::Press on all platforms for safety + /// - Batching of rapid keypresses will be implemented in Phase 2 + /// - Currently returns individual key events + /// + /// # Graceful Shutdown + /// - SIGINT (Ctrl+C) triggers shutdown and returns last event + /// - EventStream close (stdin EOF) triggers shutdown + /// - Shutdown flag can be checked after this returns + /// + /// # Errors + /// - Returns error if terminal event stream encounters an error + /// - EventStream close is handled gracefully as shutdown signal + /// + /// # Example + /// ```no_run + /// # use atuin_ai::tui::EventLoop; + /// # async fn example() -> eyre::Result<()> { + /// let mut event_loop = EventLoop::new(); + /// while !event_loop.is_shutdown() { + /// match event_loop.run().await? { + /// // Handle events... + /// # _ => break, + /// } + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn run(&mut self) -> Result { + // Create async event stream for keyboard/terminal events + let mut reader = EventStream::new(); + + // Get or create the tick timer (reused across calls to maintain timing) + // Uses fast base tick for responsive streaming; spinner timing handled in AppState + let tick_timer = self.tick_timer.get_or_insert_with(|| { + let mut interval = time::interval(BASE_TICK_INTERVAL); + // Skip the first immediate tick + interval.reset(); + interval + }); + + loop { + if self.shutdown { + break; + } + + // Biased select: prioritize stream > keyboard > tick + let event = tokio::select! { + biased; + + // Priority 1: Stream data (placeholder for Phase 3) + // In Phase 3, this will be: + // Some(chunk) = stream_rx.recv() => { ... } + + // Priority 2: Keyboard input + maybe_event = reader.next() => { + match maybe_event { + Some(Ok(Event::Key(key))) => { + // Filter to Press events only for cross-platform safety + if key.kind == KeyEventKind::Press { + // Note: Rapid keypress batching will be implemented in Phase 2 + // when we integrate with the state machine. + // For now, just return individual key events. + Some(AppEvent::Key(key)) + } else { + None + } + } + Some(Ok(Event::Resize(w, h))) => { + Some(AppEvent::Resize(w, h)) + } + Some(Err(e)) => { + return Err(eyre!("terminal event error: {}", e)); + } + None => { + // EventStream closed (stdin EOF) - trigger shutdown + self.shutdown = true; + None + } + _ => { + // Ignore other event types (mouse, focus, etc.) + None + } + } + } + + // Priority 3: Tick for spinner animation + _ = tick_timer.tick() => { + Some(AppEvent::Tick) + } + + // SIGINT handling (Ctrl+C) - cross-platform + _ = tokio::signal::ctrl_c() => { + self.shutdown = true; + // Return one more event to allow graceful shutdown handling + Some(AppEvent::Tick) + } + }; + + if let Some(app_event) = event { + return Ok(app_event); + } + } + + // Loop exited due to shutdown - return final tick to allow cleanup + Ok(AppEvent::Tick) + } + + /// Check if the event loop has been signaled to shut down. + /// + /// This can be used to cleanly exit the main TUI loop after receiving + /// a shutdown signal (Ctrl+C, stdin close, etc.) + pub fn is_shutdown(&self) -> bool { + self.shutdown + } + + /// Signal the event loop to shut down. + /// + /// The shutdown will take effect on the next iteration of `run()`. + pub fn shutdown(&mut self) { + self.shutdown = true; + } + + /// Poll for next event and apply to app state. + /// + /// This is a convenience method that combines `run()` with `App` state updates. + /// Returns true if app should continue, false if should exit. + /// + /// # Example + /// ```no_run + /// # use atuin_ai::tui::{EventLoop, App}; + /// # async fn example() -> eyre::Result<()> { + /// let mut event_loop = EventLoop::new(); + /// let mut app = App::new(); + /// + /// while event_loop.poll_and_apply(&mut app).await? { + /// // Render app state... + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn poll_and_apply(&mut self, app: &mut App) -> Result { + let event = self.run().await?; + + match event { + AppEvent::Key(key) => { + app.handle_key(key); + } + AppEvent::Tick => { + app.state.tick(); + } + AppEvent::Resize(_, _) => { + // Render will be triggered anyway + } + AppEvent::StreamChunk(_) | AppEvent::StreamDone | AppEvent::StreamError(_) => { + // Placeholder for Phase 3 + } + } + + Ok(!app.state.should_exit) + } +} + +impl Default for EventLoop { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_event_loop_creation() { + let event_loop = EventLoop::new(); + assert!(!event_loop.shutdown); + } + + #[test] + fn test_shutdown_flag() { + let mut event_loop = EventLoop::new(); + assert!(!event_loop.is_shutdown()); + + event_loop.shutdown(); + assert!(event_loop.is_shutdown()); + } + + // Note: Cannot easily test run() in unit tests since it requires a TTY. + // Integration tests should verify: + // 1. Tick events are generated at 150ms intervals + // 2. Keyboard events are properly filtered to Press only + // 3. Rapid keypresses are batched + // 4. SIGINT triggers graceful shutdown + // 5. Resize events are propagated correctly +} diff --git a/crates/atuin-ai/src/tui/mod.rs b/crates/atuin-ai/src/tui/mod.rs new file mode 100644 index 00000000..dbf4457b --- /dev/null +++ b/crates/atuin-ai/src/tui/mod.rs @@ -0,0 +1,14 @@ +pub mod app; +pub mod event; +pub mod render; +pub mod spinner; +pub mod state; +pub mod terminal; +pub mod view_model; + +pub use app::App; +pub use event::{AppEvent, EventLoop}; +pub use render::{RenderContext, calculate_needed_height, markdown_to_spans}; +pub use state::{AppMode, AppState, ConversationEvent, ExitAction}; +pub use terminal::{TerminalGuard, install_panic_hook}; +pub use view_model::{Block, Blocks, Content}; diff --git a/crates/atuin-ai/src/tui/render.rs b/crates/atuin-ai/src/tui/render.rs new file mode 100644 index 00000000..0b6341e6 --- /dev/null +++ b/crates/atuin-ai/src/tui/render.rs @@ -0,0 +1,674 @@ +use atuin_client::theme::{Meaning, Theme}; +use pulldown_cmark::{Event, Parser, Tag, TagEnd}; +use ratatui::{ + Frame, + backend::FromCrossterm, + layout::{Alignment, Constraint, Direction, Layout, Rect}, + style::{Modifier, Style}, + text::{Line, Span}, + widgets::{Block as RatatuiBlock, Borders, Padding, Paragraph, Wrap}, +}; +use tui_textarea::TextArea; + +use super::spinner::active_frame; +use super::state::AppState; +use super::view_model::{Blocks, Content, WarningKind}; + +/// Fixed card width for the TUI +const CARD_WIDTH: u16 = 64; + +pub struct RenderContext<'a> { + pub theme: &'a Theme, + pub anchor_col: u16, + pub textarea: Option<&'a TextArea<'static>>, + /// Maximum viewport height (for scroll calculations) + pub max_height: u16, +} + +/// Calculate the height needed to render the current state. +/// Used to dynamically resize the viewport before rendering. +pub fn calculate_needed_height(state: &AppState) -> u16 { + use super::state::AppMode; + + let view = Blocks::from_state(state); + let content_width = usize::from(CARD_WIDTH.saturating_sub(4)).max(1); + + let mut total_height = 0u16; + for (idx, block) in view.items.iter().enumerate() { + if idx > 0 { + total_height = total_height.saturating_add(1); // separator + total_height = total_height.saturating_add(1); // leading blank after separator + } + total_height = + total_height.saturating_add(calculate_block_height(&block.content, content_width)); + } + + // In Streaming/Generating mode, always reserve space for spinner block even during + // the 200ms delay when it's not yet shown. This prevents the UI from briefly + // shrinking and scrolling away the user message. + let has_spinner_block = view.items.iter().any(|b| { + b.content + .iter() + .any(|c| matches!(c, Content::Spinner { .. })) + }); + if matches!(state.mode, AppMode::Streaming | AppMode::Generating) && !has_spinner_block { + // Reserve space for separator (2 lines) + spinner block (1 line) + total_height = total_height.saturating_add(3); + } + + // Add borders (2) + top padding (1), minimum 5 + total_height.saturating_add(3).max(5) +} + +/// Main render function: derives view model from state, then renders it +pub fn render(frame: &mut Frame, state: &AppState, ctx: &RenderContext) { + // PURE DERIVATION: view model is always rebuilt from state + let view = Blocks::from_state(state); + + // Render the derived view model + render_view(frame, &view, ctx); +} + +fn render_view(frame: &mut Frame, view: &Blocks, ctx: &RenderContext) { + let area = frame.area(); + + // Calculate frame dimensions (fixed width, min 32 if terminal is narrow) + let desired_width = CARD_WIDTH.min(area.width.saturating_sub(2)).max(32); + let content_width = usize::from(desired_width.saturating_sub(4)).max(1); + + // Position at anchor_col + let max_x = area.x + area.width.saturating_sub(desired_width); + let preferred_x = area.x + ctx.anchor_col.saturating_sub(2); + + // Calculate height from view model + let mut total_height = 0u16; + for (idx, block) in view.items.iter().enumerate() { + if idx > 0 { + total_height = total_height.saturating_add(1); // separator + total_height = total_height.saturating_add(1); // leading blank after separator + } + total_height = + total_height.saturating_add(calculate_block_height(&block.content, content_width)); + } + + let desired_height = total_height + .saturating_add(3) // borders (2) + top padding (1), no bottom padding + .max(5); + + // Cap card height at viewport height to prevent overflow + let actual_height = desired_height.min(area.height); + + // Calculate scroll offset (scroll to show bottom content when overflowing) + let scroll_offset = desired_height.saturating_sub(actual_height); + + let card = Rect { + x: preferred_x.min(max_x), + y: area.y, + width: desired_width, + height: actual_height, + }; + + // Get title from first block (if any) + let title = view + .items + .first() + .and_then(|b| b.title.as_deref()) + .unwrap_or("Describe the command you'd like to generate:"); + + // Create bordered frame + // Padding: left=1, right=1, top=1, bottom=0 (blocks have trailing blanks) + let outer_block = RatatuiBlock::default() + .borders(Borders::ALL) + .title(title) + .title_bottom(Line::from(view.footer).alignment(Alignment::Right)) + .padding(Padding::new(1, 1, 1, 0)); + + let inner_area = outer_block.inner(card); + frame.render_widget(outer_block, card); + + // Render blocks (with scroll offset for overflowing content) + render_blocks_content(frame, view, ctx, inner_area, card.width, scroll_offset); +} + +fn render_blocks_content( + frame: &mut Frame, + view: &Blocks, + ctx: &RenderContext, + area: Rect, + card_width: u16, + scroll_offset: u16, +) { + let content_width = usize::from(area.width).max(1); + + // Build layout constraints for full content + let mut constraints = Vec::new(); + let mut block_heights = Vec::new(); + for (idx, block) in view.items.iter().enumerate() { + if idx > 0 { + constraints.push(Constraint::Length(1)); // separator + constraints.push(Constraint::Length(1)); // leading blank after separator + block_heights.push(1); + block_heights.push(1); + } + let height = calculate_block_height(&block.content, content_width); + constraints.push(Constraint::Length(height)); + block_heights.push(height); + } + + if constraints.is_empty() { + return; + } + + // Calculate cumulative heights to find which blocks are visible after scrolling + let mut cumulative: Vec = Vec::with_capacity(block_heights.len() + 1); + cumulative.push(0); + for h in &block_heights { + cumulative.push(cumulative.last().unwrap() + h); + } + + // Render each chunk, offsetting by scroll_offset and clipping to visible area + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints(constraints) + .split(area); + + let mut chunk_idx = 0; + for (idx, block) in view.items.iter().enumerate() { + if idx > 0 { + // Check if separator is visible (its position minus scroll_offset) + let sep_start = cumulative[chunk_idx]; + if sep_start >= scroll_offset && sep_start < scroll_offset + area.height { + let adjusted_chunk = Rect { + y: area.y + sep_start - scroll_offset, + ..chunks[chunk_idx] + }; + render_separator(frame, adjusted_chunk, ctx, card_width); + } + chunk_idx += 1; + chunk_idx += 1; // skip leading blank + } + + // Check if this block is at least partially visible + let block_start = cumulative[chunk_idx]; + let block_end = cumulative[chunk_idx + 1]; + + // Block is visible if it starts before viewport end and ends after viewport start + if block_start < scroll_offset + area.height && block_end > scroll_offset { + // Calculate visible portion + let visible_start = block_start.max(scroll_offset); + let visible_end = block_end.min(scroll_offset + area.height); + + let adjusted_chunk = Rect { + x: area.x, + y: area.y + visible_start - scroll_offset, + width: area.width, + height: visible_end - visible_start, + }; + + render_block_content(frame, &block.content, adjusted_chunk, ctx); + } + + chunk_idx += 1; + } +} + +/// Render all content items in a block +fn render_block_content(frame: &mut Frame, content: &[Content], area: Rect, ctx: &RenderContext) { + if content.is_empty() { + return; + } + + let content_width = usize::from(area.width).max(1); + + // Build layout constraints for each content item WITH spacing between items + let mut constraints = Vec::new(); + for (idx, c) in content.iter().enumerate() { + if idx > 0 { + constraints.push(Constraint::Length(1)); // blank line between items + } + constraints.push(Constraint::Length(calculate_single_content_height( + c, + content_width, + ))); + } + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints(constraints) + .split(area); + + let mut chunk_idx = 0; + for (idx, item) in content.iter().enumerate() { + if idx > 0 { + chunk_idx += 1; // skip the blank line chunk + } + render_single_content(frame, item, chunks[chunk_idx], ctx); + chunk_idx += 1; + } +} + +/// Render a single content item using ratatui's native wrapping. +/// Symbol is rendered at column 0, text wraps in columns 2+ (offset area). +fn render_single_content(frame: &mut Frame, content: &Content, area: Rect, ctx: &RenderContext) { + // Helper to create offset text area (2 chars for symbol column) + let text_area = Rect { + x: area.x.saturating_add(2), + y: area.y, + width: area.width.saturating_sub(2), + height: area.height, + }; + + match content { + Content::Input { text, active, .. } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Guidance)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + // Render ">" symbol at column 0 + render_symbol(frame, ">", symbol_style, area); + + if *active { + // Active input: render TextArea widget (handles cursor display) + if let Some(textarea) = ctx.textarea { + frame.render_widget(textarea, text_area); + } + } else { + // Inactive input: render as plain paragraph + let paragraph = Paragraph::new(text.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + } + + Content::Command { text, faded } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Important)); + let mut text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + if *faded { + text_style = text_style.add_modifier(Modifier::DIM); + } + + render_symbol(frame, "$", symbol_style, area); + + let paragraph = Paragraph::new(text.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Text { markdown } => { + // No symbol, just indent - render directly in offset area + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + let paragraph = Paragraph::new(markdown.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Error { message } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::AlertError)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + render_symbol(frame, "!", symbol_style, area); + + let paragraph = Paragraph::new(message.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Warning { + kind, + text, + pending_confirm, + } => { + let (symbol, meaning) = match kind { + WarningKind::Danger => ("!", Meaning::AlertError), + WarningKind::LowConfidence => ("?", Meaning::AlertWarn), + }; + let symbol_style = Style::from_crossterm(ctx.theme.as_style(meaning)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + let display_text = if *pending_confirm { + "Press Enter again to run this dangerous command" + } else { + text.as_str() + }; + + render_symbol(frame, symbol, symbol_style, area); + + let paragraph = Paragraph::new(display_text) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Spinner { + frame: spinner_frame, + status_text, + } => { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Annotation)); + let symbol = active_frame(*spinner_frame); + + render_symbol(frame, symbol, style, area); + + let paragraph = Paragraph::new(status_text.as_str()).style(style); + frame.render_widget(paragraph, text_area); + } + + Content::ToolStatus { + completed_count, + current_label, + frame: spinner_frame, + } => { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Annotation)); + + let (symbol, text) = if let Some(label) = current_label { + let spinner = active_frame(*spinner_frame); + let text = if *completed_count > 0 { + format!( + "{} (used {} tool{})", + label, + completed_count, + if *completed_count == 1 { "" } else { "s" } + ) + } else { + label.clone() + }; + (spinner, text) + } else { + ( + "\u{2713}", + format!( + "Used {} tool{}", + completed_count, + if *completed_count == 1 { "" } else { "s" } + ), + ) + }; + + render_symbol(frame, symbol, style, area); + + let paragraph = Paragraph::new(text).style(style); + frame.render_widget(paragraph, text_area); + } + } +} + +/// Render a single-character symbol at the start of an area +fn render_symbol(frame: &mut Frame, symbol: &str, style: Style, area: Rect) { + let symbol_area = Rect { + x: area.x, + y: area.y, + width: 1, + height: 1, + }; + frame.render_widget(Paragraph::new(symbol).style(style), symbol_area); +} + +fn render_separator(frame: &mut Frame, area: Rect, ctx: &RenderContext, card_width: u16) { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Muted)); + + // Build separator: ├ + ─ repeated + ┤ spanning the full card width + // -2 for the ├ and ┤ characters themselves + let inner_width = card_width.saturating_sub(2) as usize; + let separator = format!( + "\u{251c}{}\u{2524}", // ├ ... ┤ + "\u{2500}".repeat(inner_width) // ─ + ); + + let paragraph = Paragraph::new(Span::styled(separator, style)); + + // Render at x offset to overlap the border (area is inside padding, border is 2 chars left) + let sep_area = Rect { + x: area.x.saturating_sub(2), // move left to overlap left border + y: area.y, + width: card_width, + height: 1, + }; + frame.render_widget(paragraph, sep_area); +} + +/// Calculate total height for all content items in a block +fn calculate_block_height(content: &[Content], width: usize) -> u16 { + let content_height: u16 = content + .iter() + .map(|c| calculate_single_content_height(c, width)) + .sum(); + + // Add spacing between items (n-1 blank lines for n items) + let spacing = if content.len() > 1 { + (content.len() - 1) as u16 + } else { + 0 + }; + + // Add 1 for trailing blank line (padding after content) + content_height.saturating_add(spacing).saturating_add(1) +} + +/// Calculate height for a single content item. +/// Uses ratatui's Paragraph::line_count for consistency with rendering. +fn calculate_single_content_height(content: &Content, width: usize) -> u16 { + // Text area is offset by 2 for symbol column + let text_width = width.saturating_sub(2); + + match content { + // Input uses word wrapping (WrapMode::Word) in TextArea, which can produce + // more lines than character wrapping since it won't break words mid-word + Content::Input { text, active, .. } => { + if *active { + // For active input, use word-wrap line counting to match TextArea behavior + let (lines, last_line_width) = + word_wrap_line_count_with_last_width(text, text_width); + // Only add extra line for cursor if the last line is full + if last_line_width >= text_width { + lines.saturating_add(1) + } else { + lines + } + } else { + line_count_wrapped(text, text_width) + } + } + Content::Command { text, .. } => line_count_wrapped(text, text_width), + Content::Text { markdown } => line_count_wrapped(markdown, text_width), + Content::Error { message } => line_count_wrapped(message, text_width), + Content::Warning { + text, + pending_confirm, + .. + } => { + let display_text = if *pending_confirm { + "Press Enter again to run this dangerous command" + } else { + text.as_str() + }; + line_count_wrapped(display_text, text_width) + } + Content::Spinner { .. } => 1, + Content::ToolStatus { .. } => 1, + } +} + +/// Count lines when text is wrapped at given width. +/// Uses ratatui's Paragraph::line_count for accurate wrapping calculation. +fn line_count_wrapped(text: &str, width: usize) -> u16 { + if width == 0 { + return 1; + } + + let paragraph = Paragraph::new(text).wrap(Wrap { trim: false }); + paragraph.line_count(width as u16).max(1) as u16 +} + +/// Count lines using word-wrap algorithm (matches TextArea's WrapMode::Word). +/// Words won't be broken mid-word, so this may produce more lines than character wrapping. +/// Returns (line_count, last_line_width) so caller can determine if cursor needs extra space. +fn word_wrap_line_count_with_last_width(text: &str, width: usize) -> (u16, usize) { + if width == 0 || text.is_empty() { + return (1, 0); + } + + let mut line_count = 0u16; + let mut current_line_width = 0usize; + + for line in text.lines() { + if line.is_empty() { + line_count += 1; + current_line_width = 0; + continue; + } + + let mut line_started = false; + + for word in line.split_whitespace() { + let word_width = unicode_width::UnicodeWidthStr::width(word); + + if !line_started { + // First word on line + if word_width > width { + // Word is longer than width, it will be split by character + // Count how many lines it takes + line_count += word_width.div_ceil(width) as u16; + current_line_width = word_width % width; + if current_line_width == 0 { + current_line_width = 0; + line_started = false; + } else { + line_started = true; + } + } else { + current_line_width = word_width; + line_started = true; + } + } else { + // Subsequent word - need space before it + let needed = current_line_width + 1 + word_width; + if needed > width { + // Word doesn't fit, start new line + line_count += 1; + if word_width > width { + // Word itself is too long, will be split + line_count += word_width.div_ceil(width) as u16; + current_line_width = word_width % width; + if current_line_width == 0 { + line_started = false; + } + } else { + current_line_width = word_width; + } + } else { + current_line_width = needed; + } + } + } + + // Count the last line of this logical line + if line_started { + line_count += 1; + } + } + + // Handle case where text has no lines() output (empty or just whitespace) + if line_count == 0 { + line_count = 1; + current_line_width = 0; + } + + (line_count, current_line_width) +} + +/// Convert markdown to styled spans (existing function, kept as-is) +pub fn markdown_to_spans<'a>(text: &'a str, theme: &'a Theme) -> Vec> { + let parser = Parser::new(text); + let mut lines: Vec>> = vec![Vec::new()]; + let mut current_line = 0; + + let base_style = Style::from_crossterm(theme.as_style(Meaning::Base)); + let code_style = Style::from_crossterm(theme.as_style(Meaning::Important)); + let mut style_stack: Vec