use chrono::Local;
use glob::glob;
use regex::Regex;
use rustyline::DefaultEditor;
use rustyline::error::ReadlineError;
use serde::{Deserialize, Serialize};
use std::env;
use std::fs;
use std::io::{self, Write};
use std::path::PathBuf;
use std::process::Command;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use std::thread;
use std::time::Duration;
use walkdir::WalkDir;
use menta::{GenerateTextRequest, ModelMessage, Tool, ToolChoice, ToolExecute, generate_text};
const RESET: &str = "\x1b[0m";
const DIM: &str = "\x1b[2m";
const BLUE: &str = "\x1b[34m";
const CYAN: &str = "\x1b[36m";
const GREEN: &str = "\x1b[32m";
const YELLOW: &str = "\x1b[33m";
const RED: &str = "\x1b[31m";
fn paint(color: &str, value: impl AsRef<str>) -> String {
format!("{color}{}{RESET}", value.as_ref())
}
struct SpinnerState {
active: AtomicBool,
paused: AtomicBool,
}
struct LoadingIndicator {
state: Arc<SpinnerState>,
handle: Option<thread::JoinHandle<()>>,
}
impl LoadingIndicator {
fn start(message: &'static str) -> Self {
let state = Arc::new(SpinnerState {
active: AtomicBool::new(true),
paused: AtomicBool::new(false),
});
let state_for_thread = Arc::clone(&state);
let handle = thread::spawn(move || {
let frames = ['|', '/', '-', '\\'];
let mut index = 0;
while state_for_thread.active.load(Ordering::Relaxed) {
if state_for_thread.paused.load(Ordering::Relaxed) {
thread::sleep(Duration::from_millis(25));
continue;
}
print!(
"\r{} {}",
paint(CYAN, frames[index % frames.len()].to_string()),
paint(DIM, message),
);
let _ = io::stdout().flush();
index += 1;
thread::sleep(Duration::from_millis(100));
}
});
Self {
state,
handle: Some(handle),
}
}
fn state(&self) -> Arc<SpinnerState> {
Arc::clone(&self.state)
}
fn stop(mut self) {
self.state.active.store(false, Ordering::Relaxed);
if let Some(handle) = self.handle.take() {
let _ = handle.join();
}
print!("\r{}\r", " ".repeat(40));
let _ = io::stdout().flush();
}
}
struct SpinnerPause {
state: Arc<SpinnerState>,
}
impl SpinnerPause {
fn new(state: Arc<SpinnerState>) -> Self {
state.paused.store(true, Ordering::Relaxed);
print!("\r{}\r", " ".repeat(60));
let _ = io::stdout().flush();
Self { state }
}
}
impl Drop for SpinnerPause {
fn drop(&mut self) {
self.state.paused.store(false, Ordering::Relaxed);
}
}
fn logged_tool<T>(verbose: bool, spinner_state: Arc<SpinnerState>) -> Tool
where
T: ToolExecute,
{
let definition = T::definition();
let name = definition.name.clone();
Tool::new_async(
definition.name,
definition.description,
definition.input_schema,
definition.output_schema,
move |input| {
let name = name.clone();
let spinner_state = Arc::clone(&spinner_state);
async move {
let _pause = SpinnerPause::new(spinner_state);
println!(
"{} {} {}",
paint(YELLOW, "tool>"),
paint(YELLOW, &name),
paint(DIM, &input)
);
let parsed = serde_json::from_str::<T>(&input)
.map_err(|error| format!("invalid tool input for {name}: {error}"))?;
let output = parsed.execute().await?;
let output = serde_json::to_string(&output)
.map_err(|error| format!("invalid tool output for {name}: {error}"))?;
if verbose {
println!(
"{} {} {}",
paint(YELLOW, "tool<"),
paint(YELLOW, &name),
paint(DIM, &output)
);
}
Ok(output)
}
},
)
}
#[derive(Deserialize, Tool)]
#[tool(description = "Get the weather in a location")]
struct WeatherTool {
#[description = "The location to get the weather for"]
location: String,
}
#[derive(Serialize)]
struct WeatherOutput {
temperature: i32,
conditions: String,
}
impl ToolExecute for WeatherTool {
type Output = WeatherOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
Ok(WeatherOutput {
temperature: 72,
conditions: format!("sunny in {}", self.location),
})
}
}
#[derive(Deserialize, Tool)]
#[tool(description = "Run a local shell command")]
struct BashTool {
#[description = "The shell command to run"]
command: String,
}
#[derive(Serialize)]
struct BashOutput {
success: bool,
status: i32,
stdout: String,
stderr: String,
}
impl ToolExecute for BashTool {
type Output = BashOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
let output = Command::new("sh")
.arg("-c")
.arg(&self.command)
.output()
.map_err(|error| error.to_string())?;
Ok(BashOutput {
success: output.status.success(),
status: output.status.code().unwrap_or(-1),
stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
})
}
}
#[derive(Deserialize, Tool)]
#[tool(description = "Write content to a local file")]
struct WriteTool {
#[description = "Path to the file to write"]
path: String,
#[description = "Full file content to write"]
content: String,
}
#[derive(Serialize)]
struct WriteOutput {
path: String,
bytes_written: usize,
}
impl ToolExecute for WriteTool {
type Output = WriteOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
if let Some(parent) = PathBuf::from(&self.path).parent() {
if !parent.as_os_str().is_empty() {
fs::create_dir_all(parent).map_err(|error| error.to_string())?;
}
}
fs::write(&self.path, &self.content).map_err(|error| error.to_string())?;
Ok(WriteOutput {
path: self.path.clone(),
bytes_written: self.content.len(),
})
}
}
#[derive(Deserialize, Tool)]
#[tool(description = "Read a local file")]
struct ReadTool {
#[description = "Path to the file to read"]
path: String,
}
#[derive(Serialize)]
struct ReadOutput {
path: String,
content: String,
}
impl ToolExecute for ReadTool {
type Output = ReadOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
let content = fs::read_to_string(&self.path).map_err(|error| error.to_string())?;
Ok(ReadOutput {
path: self.path.clone(),
content,
})
}
}
#[derive(Deserialize, Tool)]
#[tool(description = "Find files by glob pattern")]
struct GlobTool {
#[description = "Glob pattern like **/*.rs or src/**/*.rs"]
pattern: String,
}
#[derive(Serialize)]
struct GlobOutput {
matches: Vec<String>,
}
impl ToolExecute for GlobTool {
type Output = GlobOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
let matches = glob(&self.pattern)
.map_err(|error| error.to_string())?
.filter_map(|entry| entry.ok())
.map(|path| path.display().to_string())
.collect::<Vec<_>>();
Ok(GlobOutput { matches })
}
}
#[derive(Deserialize, Tool)]
#[tool(description = "Search file contents with a regex pattern")]
struct GrepTool {
#[description = "Regex pattern to search for"]
pattern: String,
#[description = "Root path to search from, defaults to current directory"]
path: Option<String>,
}
#[derive(Serialize)]
struct GrepOutput {
matches: Vec<String>,
}
impl ToolExecute for GrepTool {
type Output = GrepOutput;
async fn execute(&self) -> std::result::Result<Self::Output, String> {
let regex = Regex::new(&self.pattern).map_err(|error| error.to_string())?;
let root = self.path.as_deref().unwrap_or(".");
let mut matches = Vec::new();
for entry in WalkDir::new(root) {
let entry = entry.map_err(|error| error.to_string())?;
if !entry.file_type().is_file() {
continue;
}
let path = entry.path();
let Ok(content) = fs::read_to_string(path) else {
continue;
};
for (index, line) in content.lines().enumerate() {
if regex.is_match(line) {
matches.push(format!("{}:{}:{}", path.display(), index + 1, line));
}
}
}
Ok(GrepOutput { matches })
}
}
#[tokio::main]
async fn main() {
let verbose = env::args().skip(1).any(|arg| arg == "--verbose");
println!("{}", paint(BLUE, "menta agent example"));
println!(
"{}",
paint(
DIM,
"commands: /new, /model [model-id], exit, quit, !<shell-command>"
)
);
if verbose {
println!("{}", paint(DIM, "verbose tool logging enabled"));
}
let mut model = String::from("openai/gpt-4.1-mini");
let mut history = vec![ModelMessage::system(system_prompt())];
let mut editor = DefaultEditor::new().expect("failed to initialize line editor");
loop {
let prompt = match editor.readline(&format!("[{model}] > ")) {
Ok(prompt) => prompt,
Err(ReadlineError::Interrupted) => {
println!();
continue;
}
Err(ReadlineError::Eof) => break,
Err(error) => {
eprintln!("error: failed to read input: {error}");
break;
}
};
let prompt = prompt.trim();
if prompt.is_empty() {
continue;
}
let _ = editor.add_history_entry(prompt);
if matches!(prompt, "exit" | "quit") {
break;
}
if prompt == "/new" {
history = vec![ModelMessage::system(system_prompt())];
println!("started new session");
continue;
}
if let Some(next_model) = prompt.strip_prefix("/model") {
let next_model = next_model.trim();
if next_model.is_empty() {
println!("{} {}", paint(BLUE, "current model:"), paint(CYAN, &model));
} else {
model = next_model.to_string();
println!(
"{} {}",
paint(BLUE, "switched model to"),
paint(CYAN, &model)
);
}
continue;
}
if let Some(command) = prompt.strip_prefix('!') {
let command = command.trim();
if command.is_empty() {
eprintln!("{} missing shell command after !", paint(RED, "error:"));
continue;
}
match Command::new("sh").arg("-c").arg(command).output() {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
if !stdout.trim().is_empty() {
print!("{stdout}");
}
if !stderr.trim().is_empty() {
eprint!("{stderr}");
}
if !output.status.success() {
eprintln!(
"{} command exited with status {}",
paint(RED, "error:"),
output.status
);
}
}
Err(error) => eprintln!("{} failed to run command: {error}", paint(RED, "error:")),
}
continue;
}
let mut messages = history.clone();
messages.push(ModelMessage::user(prompt));
let loading = LoadingIndicator::start("thinking...");
let spinner_state = loading.state();
let result = generate_text(
GenerateTextRequest::new()
.model(model.clone())
.messages(messages)
.tools(vec![
logged_tool::<WeatherTool>(verbose, Arc::clone(&spinner_state)),
logged_tool::<BashTool>(verbose, Arc::clone(&spinner_state)),
logged_tool::<WriteTool>(verbose, Arc::clone(&spinner_state)),
logged_tool::<ReadTool>(verbose, Arc::clone(&spinner_state)),
logged_tool::<GlobTool>(verbose, Arc::clone(&spinner_state)),
logged_tool::<GrepTool>(verbose, Arc::clone(&spinner_state)),
])
.tool_choice(ToolChoice::Auto)
.max_steps(4),
)
.await;
loading.stop();
match result {
Ok(result) => {
println!("{} {}", paint(GREEN, "assistant>"), result.output);
history.push(ModelMessage::user(prompt));
history.push(ModelMessage::assistant_text(result.text));
}
Err(error) => eprintln!("{} {error}", paint(RED, "error:")),
}
}
}
fn system_prompt() -> String {
let date = Local::now().format("%Y-%m-%d").to_string();
let pwd = env::current_dir()
.map(|path| path.display().to_string())
.unwrap_or_else(|_| String::from("unknown"));
format!(
"You are a local coding assistant. Current date: {date}. Current working directory: {pwd}. Use tools when needed: weather, bash, write, read, glob, grep. Prefer tools for filesystem and shell tasks."
)
}