LazyBoy2/crates/grokboy/src/main.rs

use anyhow::{Context, Result, anyhow};
use grokboy_core::{
    AGENT_SYSTEM, ChatMessage, Config, DEFAULT_MAX_ROUNDS, Session, ToolContext, execute_tool,
    load_or_create, run_agent, save_session, stream_chat, tool_definitions,
};
use serde_json::json;
use std::io::{self, Write};
use std::process::ExitCode;

const CHAT_SYSTEM: &str = "You are GrokBoy, a concise local coding assistant. Prefer clear, short answers. Traditional Chinese is welcome when the user writes in Chinese.";

#[tokio::main]
async fn main() -> ExitCode {
    if let Err(err) = run().await {
        eprintln!("error: {err:#}");
        return ExitCode::FAILURE;
    }
    ExitCode::SUCCESS
}

async fn run() -> Result<()> {
    let mut args = std::env::args().skip(1).collect::<Vec<_>>();
    if args.is_empty() {
        print_help();
        return Ok(());
    }
    let cmd = args.remove(0);
    match cmd.as_str() {
        "chat" => cmd_chat().await,
        "run" => cmd_run(&args).await,
        "smoke" => cmd_smoke().await,
        "version" | "-V" | "--version" => {
            println!("grokboy {}", env!("CARGO_PKG_VERSION"));
            Ok(())
        }
        "help" | "-h" | "--help" => {
            print_help();
            Ok(())
        }
        other => {
            eprintln!("unknown command: {other}\n");
            print_help();
            Err(anyhow!("unknown command"))
        }
    }
}

fn print_help() {
    println!(
        "\
GrokBoy — minimal local CLI agent (P1: tools + ReAct)

USAGE:
  grokboy chat                      Interactive streaming chat (no tools)
  grokboy run \"<prompt>\"            One-shot agent with tools
  grokboy run --session <id> \"...\"  Continue a saved session
  grokboy smoke                     Offline tool checks (no API key required)
  grokboy version
  grokboy help

ENV:
  GROKBOY_API_KEY       API key (or XAI_API_KEY / OPENAI_API_KEY)
  GROKBOY_BASE_URL      default https://api.x.ai/v1
  GROKBOY_MODEL         default grok-4.6

Sessions are stored under ~/.grokboy/sessions/<id>.json
"
    );
}

async fn cmd_chat() -> Result<()> {
    let config = Config::from_env().map_err(anyhow::Error::msg)?;
    println!(
        "GrokBoy chat  model={}  base={}",
        config.model, config.base_url
    );
    println!("Type a message. Empty line or /exit to quit.\n");

    let mut history = vec![ChatMessage::system(CHAT_SYSTEM)];
    let stdin = io::stdin();
    let mut stdout = io::stdout();

    loop {
        print!("you> ");
        stdout.flush().ok();
        let mut line = String::new();
        if stdin.read_line(&mut line).context("stdin")? == 0 {
            println!();
            break;
        }
        let input = line.trim();
        if input.is_empty() || input == "/exit" || input == "/quit" {
            break;
        }
        history.push(ChatMessage::user(input));

        print!("assistant> ");
        stdout.flush().ok();
        let reply = stream_chat(&config, &history, |delta| {
            print!("{delta}");
            let _ = io::stdout().flush();
        })
        .await?;
        println!("\n");
        history.push(ChatMessage::assistant(reply));
    }
    Ok(())
}

async fn cmd_run(args: &[String]) -> Result<()> {
    let mut session_id: Option<String> = None;
    let mut prompt_parts: Vec<String> = Vec::new();
    let mut i = 0;
    while i < args.len() {
        match args[i].as_str() {
            "--session" | "-s" => {
                i += 1;
                let id = args
                    .get(i)
                    .ok_or_else(|| anyhow!("--session requires an id"))?;
                session_id = Some(id.clone());
            }
            "--help" | "-h" => {
                println!(
                    "Usage: grokboy run [--session <id>] \"<prompt>\"\n\
                     Runs a one-shot tool-using agent. Creates a session under ~/.grokboy/sessions/."
                );
                return Ok(());
            }
            other => prompt_parts.push(other.to_string()),
        }
        i += 1;
    }
    let prompt = prompt_parts.join(" ");
    if prompt.trim().is_empty() {
        return Err(anyhow!("usage: grokboy run [--session <id>] \"<prompt>\""));
    }

    let config = Config::from_env().map_err(anyhow::Error::msg)?;
    let cwd = std::env::current_dir().context("cwd")?;
    let mut session = load_or_create(session_id.as_deref(), &cwd)?;

    // Ensure system prompt is present once at the start.
    if session.messages.is_empty() {
        session.push(ChatMessage::system(AGENT_SYSTEM));
    }
    session.cwd = cwd.clone();
    session.push(ChatMessage::user(&prompt));

    let tool_ctx = ToolContext::new(session.cwd.clone());
    let answer = run_agent(
        &config,
        &mut session.messages,
        &tool_ctx,
        DEFAULT_MAX_ROUNDS,
    )
    .await?;

    session.touch();
    let path = save_session(&session)?;
    println!("{answer}");
    eprintln!(
        "\n[session {} saved → {}]",
        session.id,
        path.display()
    );
    Ok(())
}

async fn cmd_smoke() -> Result<()> {
    println!("GrokBoy smoke (offline tools)…");
    let stamp = uuid_like();
    let dir = std::env::temp_dir().join(format!("grokboy-smoke-{stamp}"));
    std::fs::create_dir_all(&dir).context("temp dir")?;
    let ctx = ToolContext::new(dir.clone());

    // write_file
    let w = execute_tool(
        &ctx,
        "write_file",
        &json!({"path": "note.txt", "content": "smoke ok\n第二行"}).to_string(),
    )
    .await;
    let w: serde_json::Value = serde_json::from_str(&w)?;
    assert_ok(&w, "write_file")?;
    println!("  write_file  ok");

    // read_file
    let r = execute_tool(
        &ctx,
        "read_file",
        &json!({"path": "note.txt"}).to_string(),
    )
    .await;
    let r: serde_json::Value = serde_json::from_str(&r)?;
    assert_ok(&r, "read_file")?;
    if r["content"].as_str() != Some("smoke ok\n第二行") {
        return Err(anyhow!("read_file content mismatch: {r}"));
    }
    println!("  read_file   ok");

    // list_dir
    let l = execute_tool(&ctx, "list_dir", &json!({"path": "."}).to_string()).await;
    let l: serde_json::Value = serde_json::from_str(&l)?;
    assert_ok(&l, "list_dir")?;
    let names: Vec<&str> = l["entries"]
        .as_array()
        .context("entries")?
        .iter()
        .filter_map(|e| e["name"].as_str())
        .collect();
    if !names.contains(&"note.txt") {
        return Err(anyhow!("list_dir missing note.txt: {l}"));
    }
    println!("  list_dir    ok");

    // shell
    let s = execute_tool(
        &ctx,
        "shell",
        &json!({"command": "printf 'hi'; wc -c < note.txt"}).to_string(),
    )
    .await;
    let s: serde_json::Value = serde_json::from_str(&s)?;
    assert_ok(&s, "shell")?;
    if s["exit_code"] != 0 {
        return Err(anyhow!("shell exit_code != 0: {s}"));
    }
    println!("  shell       ok");

    // path traversal should fail
    let bad = execute_tool(
        &ctx,
        "read_file",
        &json!({"path": "../outside.txt"}).to_string(),
    )
    .await;
    let bad: serde_json::Value = serde_json::from_str(&bad)?;
    if bad.get("error").is_none() {
        return Err(anyhow!("expected path traversal error, got {bad}"));
    }
    println!("  sandbox     ok");

    // tool definitions present
    let defs = tool_definitions();
    if defs.as_array().map(|a| a.len()).unwrap_or(0) != 4 {
        return Err(anyhow!("expected 4 tool defs"));
    }
    println!("  tool defs   ok");

    // session roundtrip in temp (does not require ~/.grokboy for this check —
    // we still exercise Session serialize via save into temp using core types)
    let mut sess = Session::new(dir.clone());
    sess.push(ChatMessage::system(AGENT_SYSTEM));
    sess.push(ChatMessage::user("smoke"));
    let sess_path = dir.join("session.json");
    let data = serde_json::to_vec_pretty(&sess)?;
    std::fs::write(&sess_path, data)?;
    let loaded: Session = serde_json::from_slice(&std::fs::read(&sess_path)?)?;
    if loaded.messages.len() != 2 {
        return Err(anyhow!("session roundtrip failed"));
    }
    println!("  session     ok");

    let _ = std::fs::remove_dir_all(&dir);

    // Optional live ping if API key is present (does not fail smoke).
    if Config::from_env().is_ok() {
        println!("  (API key present — skipping live call in smoke; use `run` to exercise)");
    } else {
        println!("  (no API key — live agent not checked; offline smoke passed)");
    }

    println!("smoke passed");
    Ok(())
}

fn assert_ok(v: &serde_json::Value, label: &str) -> Result<()> {
    if let Some(err) = v.get("error") {
        return Err(anyhow!("{label} failed: {err}"));
    }
    Ok(())
}

fn uuid_like() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let n = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    format!("{n}")
}