diff options
Diffstat (limited to 'crates/atuin-ai/src/user_context/interpolate.rs')
| -rw-r--r-- | crates/atuin-ai/src/user_context/interpolate.rs | 279 |
1 files changed, 0 insertions, 279 deletions
diff --git a/crates/atuin-ai/src/user_context/interpolate.rs b/crates/atuin-ai/src/user_context/interpolate.rs deleted file mode 100644 index 91e34ab4..00000000 --- a/crates/atuin-ai/src/user_context/interpolate.rs +++ /dev/null @@ -1,279 +0,0 @@ -//! Parse `.atuin/ai-context.md` files and execute embedded commands. -//! -//! Two interpolation syntaxes are supported: -//! -//! **Inline:** `!`command`` — the `!` immediately before a code span triggers -//! execution. The entire `!`...`` span is replaced with the command's stdout. -//! -//! **Block:** -//! ````markdown -//! ```! -//! command -//! ``` -//! ```` -//! A fenced code block with `!` as the info string. The block body is executed -//! as a script and the entire fenced block is replaced with stdout. -//! -//! Regular code spans and fenced code blocks (without `!`) are left untouched. - -use std::ops::Range; -use std::time::Duration; - -use pulldown_cmark::{CodeBlockKind, Event, Options, Parser, Tag, TagEnd}; - -/// A command to execute, with its byte range in the source for replacement. -#[derive(Debug)] -struct Command { - /// Byte range in the source to replace (includes the `!` for inline, or - /// the full ``` fence for blocks). - range: Range<usize>, - /// The command string to execute. - body: String, -} - -/// Maximum time for a single command. -const COMMAND_TIMEOUT: Duration = Duration::from_secs(5); - -/// Maximum bytes of stdout to capture from a single command. -const MAX_OUTPUT_BYTES: usize = 64 * 1024; - -/// Parse a context file for interpolation commands. -fn parse_commands(source: &str) -> Vec<Command> { - let parser = Parser::new_ext(source, Options::empty()); - let mut commands = Vec::new(); - - // Block state: accumulate text across multiple Text events, finalize on End. - let mut block_start: Option<usize> = None; - let mut block_body = String::new(); - - for (event, range) in parser.into_offset_iter() { - match event { - // Inline: !`command` - Event::Code(code) if range.start > 0 && source.as_bytes()[range.start - 1] == b'!' => { - commands.push(Command { - range: (range.start - 1)..range.end, - body: code.to_string(), - }); - } - - // Block: ```! ... ``` - Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(info))) if info.as_ref() == "!" => { - block_start = Some(range.start); - block_body.clear(); - } - Event::Text(text) if block_start.is_some() => { - block_body.push_str(&text); - } - Event::End(TagEnd::CodeBlock) if block_start.is_some() => { - let start = block_start.take().unwrap(); - let trimmed = block_body.trim(); - if !trimmed.is_empty() { - commands.push(Command { - range: start..range.end, - body: trimmed.to_string(), - }); - } - block_body.clear(); - } - - _ => {} - } - } - - commands -} - -/// Execute all commands in a context file and return the interpolated content. -/// -/// Commands are executed in parallel. Failed commands are replaced with an -/// error marker so the AI has visibility into what went wrong. -pub(crate) async fn interpolate(source: &str, shell: &str) -> String { - let commands = parse_commands(source); - if commands.is_empty() { - return source.to_string(); - } - - // Execute all commands in parallel. - let mut handles = Vec::with_capacity(commands.len()); - for cmd in &commands { - let shell = shell.to_string(); - let body = cmd.body.clone(); - handles.push(tokio::spawn( - async move { run_command(&shell, &body).await }, - )); - } - - // Collect results. - let mut results = Vec::with_capacity(handles.len()); - for handle in handles { - let output = match handle.await { - Ok(output) => output, - Err(e) => format!("[error: task panicked: {e}]"), - }; - results.push(output); - } - - // Rebuild the source, replacing command ranges with their output. - // Commands are in source order from the parser, but let's sort to be safe. - let mut replacements: Vec<(Range<usize>, &str)> = commands - .iter() - .zip(results.iter()) - .map(|(cmd, output)| (cmd.range.clone(), output.as_str())) - .collect(); - replacements.sort_by_key(|(range, _)| range.start); - - let mut out = String::with_capacity(source.len()); - let mut cursor = 0; - for (range, output) in &replacements { - out.push_str(&source[cursor..range.start]); - out.push_str(output); - cursor = range.end; - } - out.push_str(&source[cursor..]); - - out -} - -async fn run_command(shell: &str, body: &str) -> String { - let result = tokio::time::timeout( - COMMAND_TIMEOUT, - tokio::process::Command::new(shell) - .arg("-c") - .arg(body) - .output(), - ) - .await; - - match result { - Ok(Ok(output)) => { - if output.status.success() { - if output.stdout.len() > MAX_OUTPUT_BYTES { - let truncated = String::from_utf8_lossy(&output.stdout[..MAX_OUTPUT_BYTES]); - format!( - "{}\n[output truncated at {}KB]", - truncated.trim(), - MAX_OUTPUT_BYTES / 1024 - ) - } else { - String::from_utf8_lossy(&output.stdout).trim().to_string() - } - } else { - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - let code = output.status.code().unwrap_or(-1); - format!("[error: exit code {code}: {stderr}]") - } - } - Ok(Err(e)) => format!("[error: {e}]"), - Err(_) => format!( - "[error: command timed out after {}s]", - COMMAND_TIMEOUT.as_secs() - ), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parse_inline_command() { - let source = "Branch: !`git branch --show-current`"; - let cmds = parse_commands(source); - assert_eq!(cmds.len(), 1); - assert_eq!(cmds[0].body, "git branch --show-current"); - assert_eq!( - &source[cmds[0].range.clone()], - "!`git branch --show-current`" - ); - } - - #[test] - fn parse_inline_double_backtick() { - let source = r#"Host: !``echo `hostname` ``"#; - let cmds = parse_commands(source); - assert_eq!(cmds.len(), 1); - assert_eq!(cmds[0].body, "echo `hostname` "); - } - - #[test] - fn parse_block_command() { - let source = "Before\n\n```!\necho hello\npython3 --version\n```\n\nAfter"; - let cmds = parse_commands(source); - assert_eq!(cmds.len(), 1); - assert_eq!(cmds[0].body, "echo hello\npython3 --version"); - } - - #[test] - fn regular_code_not_matched() { - let source = "Normal `code span` and ```bash\necho hi\n```"; - let cmds = parse_commands(source); - assert_eq!(cmds.len(), 0); - } - - #[test] - fn bang_not_adjacent_not_matched() { - let source = "Exclaim! Then `code` here."; - let cmds = parse_commands(source); - // The `!` and backtick are separated by " Then ", not adjacent. - assert_eq!(cmds.len(), 0); - } - - #[test] - fn mixed_content() { - let source = "\ -# Project Context - -Branch: !`git branch --show-current` - -Regular code: `not a command` - -```! -echo $VIRTUAL_ENV -``` - -```bash -echo not interpolated -``` - -End."; - let cmds = parse_commands(source); - assert_eq!(cmds.len(), 2); - assert_eq!(cmds[0].body, "git branch --show-current"); - assert_eq!(cmds[1].body, "echo $VIRTUAL_ENV"); - } - - #[tokio::test] - async fn interpolate_replaces_inline_command() { - let source = "Branch: !`echo main`"; - let result = interpolate(source, "sh").await; - assert_eq!(result, "Branch: main"); - } - - #[tokio::test] - async fn interpolate_replaces_block_command() { - let source = "Before\n\n```!\necho hello world\n```\n\nAfter"; - let result = interpolate(source, "sh").await; - assert_eq!(result, "Before\n\nhello world\n\nAfter"); - } - - #[tokio::test] - async fn interpolate_preserves_non_command_content() { - let source = "Just plain markdown with `code` and no bangs."; - let result = interpolate(source, "sh").await; - assert_eq!(result, source); - } - - #[tokio::test] - async fn interpolate_failed_command_shows_error() { - let source = "Result: !`exit 1`"; - let result = interpolate(source, "sh").await; - assert!(result.starts_with("Result: [error:")); - } - - #[tokio::test] - async fn interpolate_multiple_commands() { - let source = "A: !`echo one` B: !`echo two`"; - let result = interpolate(source, "sh").await; - assert_eq!(result, "A: one B: two"); - } -} |