LazyBoy2/crates/grokboy/src/main.rs

use anyhow::{Context, Result, anyhow};
use grokboy_core::{
    AGENT_SYSTEM, AgentVerdict, ChatMessage, Config, DEFAULT_MAX_ROUNDS, Session, ToolContext,
    execute_tool, is_completion_tool, load_or_create, messages_char_len, run_agent, save_session,
    stream_chat, tool_definitions, truncate_messages,
};
use serde_json::json;
use std::io::{self, Write};
use std::path::PathBuf;
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 (P3: optional Playwright DOM browser)

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 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
  GROKBOY_CONTEXT_CHARS    context budget (default 100000)

Tools: shell, list_dir, read_file, write_file, report_done, report_blocked,
       browser_navigate, browser_snapshot, browser_click, browser_type, browser_eval
Sessions: ~/.grokboy/sessions/<id>.json
Browser (optional): cd tools/playwright && npm i && npx playwright install chromium
"
    );
}

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 verdict = run_agent(
        &config,
        &mut session.messages,
        &tool_ctx,
        DEFAULT_MAX_ROUNDS,
    )
    .await?;

    if let Some(url) = tool_ctx.last_browser_url_value() {
        session.last_browser_url = Some(url);
    }
    session.touch();
    let path = save_session(&session)?;
    println!("{}", verdict.message());
    eprintln!(
        "\n[verdict: {} | session {} → {}]",
        verdict.kind(),
        session.id,
        path.display()
    );
    // Non-zero exit on blocked so scripts can detect fail-closed.
    if matches!(verdict, AgentVerdict::Blocked(_)) {
        return Err(anyhow!("agent blocked: {}", verdict.message()));
    }
    Ok(())
}

async fn cmd_smoke() -> Result<()> {
    println!("GrokBoy smoke (offline P3)…");
    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");

    // completion tools
    let done = execute_tool(
        &ctx,
        "report_done",
        &json!({"message": "smoke done"}).to_string(),
    )
    .await;
    let done: serde_json::Value = serde_json::from_str(&done)?;
    if done["status"] != "done" || done["message"] != "smoke done" {
        return Err(anyhow!("report_done mismatch: {done}"));
    }
    let blocked = execute_tool(
        &ctx,
        "report_blocked",
        &json!({"reason": "smoke blocked"}).to_string(),
    )
    .await;
    let blocked: serde_json::Value = serde_json::from_str(&blocked)?;
    if blocked["status"] != "blocked" || blocked["reason"] != "smoke blocked" {
        return Err(anyhow!("report_blocked mismatch: {blocked}"));
    }
    if !is_completion_tool("report_done") || !is_completion_tool("report_blocked") {
        return Err(anyhow!("is_completion_tool failed"));
    }
    println!("  completion     ok");

    // tool definitions present (6 core + 5 browser)
    let defs = tool_definitions();
    let n_tools = defs.as_array().map(|a| a.len()).unwrap_or(0);
    if n_tools != 11 {
        return Err(anyhow!("expected 11 tool defs, got {n_tools}"));
    }
    println!("  tool defs      ok");

    // browser helper protocol (no Chromium required)
    let repo = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
    match grokboy_core::browser_self_test(&repo) {
        Ok(_) => println!("  browser proto  ok"),
        Err(e) => {
            // Soft: smoke still passes if helper missing in weird cwd; warn.
            println!("  browser proto  skip ({e:#})");
        }
    }
    // Fail-closed install hint when navigating without ensuring chromium — oneshot missing url already covered in helper self-test.
    let nav = execute_tool(
        &ToolContext::new(repo.clone()),
        "browser_navigate",
        &json!({"url": ""}).to_string(),
    )
    .await;
    let nav_v: serde_json::Value = serde_json::from_str(&nav)?;
    // empty url should error from helper; if helper missing, also error with install hint
    if nav_v.get("error").is_none() && nav_v.get("blocked").is_none() {
        // empty string might still "succeed" oddly — accept either error or blocked
        println!("  browser fail   note: {nav_v}");
    } else {
        println!("  browser fail   ok");
    }

    // context truncation
    let mut msgs = vec![
        ChatMessage::system(AGENT_SYSTEM),
        ChatMessage::user("goal"),
        ChatMessage::tool("t1", "X".repeat(50_000)),
        ChatMessage::tool("t2", "Y".repeat(50_000)),
        ChatMessage::tool("t3", "Z".repeat(50_000)),
        ChatMessage::tool("t4", "W".repeat(50_000)),
        ChatMessage::tool("t5", "V".repeat(50_000)),
        ChatMessage::user("goal again"),
    ];
    let before = messages_char_len(&msgs);
    truncate_messages(&mut msgs, 10_000);
    let after = messages_char_len(&msgs);
    if after >= before {
        return Err(anyhow!("truncate did not shrink ({before} -> {after})"));
    }
    if msgs[0].text() != AGENT_SYSTEM {
        return Err(anyhow!("truncate dropped system prompt"));
    }
    println!("  truncation     ok");

    // session roundtrip
    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);

    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}")
}