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,
        "agent" => cmd_agent(&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 (P6: confirm-before-post + scenario playbooks)

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 agent                     Interactive multi-turn agent REPL (with tools)
  grokboy agent --session <id>      Resume an agent 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)
  GROKBOY_BROWSER_HEADED   1 = always launch Chromium headed (visible; for run/agent)
  GROKBOY_HANDOFF_AUTO     1 = auto-resume handoff (tests); abort = auto-abort
  GROKBOY_CONFIRM_AUTO     1 = auto-approve confirm (tests); abort = auto-deny
                           (falls back to HANDOFF_AUTO if unset)

Tools: shell, list_dir, read_file, write_file, report_done, report_blocked,
       request_user_confirm, browser_navigate, browser_snapshot, browser_click,
       browser_type, browser_eval, browser_handoff
Sessions: ~/.grokboy/sessions/<id>.json
Browser (optional): cd tools/playwright && npm i && npx playwright install chromium
Handoff: agent pauses on login/OTP/captcha → you fix in headed Chromium → Enter
REPL: /exit /quit leave; /session show id; empty line ignored
"
    );
}

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


/// Parse `agent` CLI flags. Returns (session_id, show_help).
fn parse_agent_args(args: &[String]) -> Result<(Option<String>, bool)> {
    let mut session_id: Option<String> = None;
    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" => return Ok((session_id, true)),
            other => {
                return Err(anyhow!(
                    "unexpected argument: {other}\nusage: grokboy agent [--session <id>]"
                ));
            }
        }
        i += 1;
    }
    Ok((session_id, false))
}

async fn cmd_agent(args: &[String]) -> Result<()> {
    let (session_id, show_help) = parse_agent_args(args)?;
    if show_help {
        println!(
            "Usage: grokboy agent [--session <id>]\n\
             Interactive multi-turn ReAct agent with tools.\n\
             Creates/resumes a session under ~/.grokboy/sessions/.\n\
             Commands: /exit /quit  leave;  /session  show id;  empty line ignored.\n\
             Tip: GROKBOY_BROWSER_HEADED=1 for visible Chromium (handoff)."
        );
        return Ok(());
    }

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

    if session.messages.is_empty() {
        session.push(ChatMessage::system(AGENT_SYSTEM));
    }
    session.cwd = cwd.clone();

    // One ToolContext for the whole REPL so browser URL/state carries across turns.
    let tool_ctx = ToolContext::new(session.cwd.clone());
    if let Some(url) = &session.last_browser_url {
        if let Ok(mut g) = tool_ctx.last_browser_url.lock() {
            *g = Some(url.clone());
        }
    }

    // Persist early so the printed session id is on disk.
    let path = save_session(&session)?;

    println!(
        "GrokBoy agent  model={}  base={}",
        config.model, config.base_url
    );
    println!("session {}", session.id);
    if created_new {
        eprintln!("[new session → {}]", path.display());
    } else {
        eprintln!("[resumed → {}]", path.display());
    }
    println!(
        "互動式多輪代理（含工具）。輸入訊息後會跑 ReAct；/exit 或 /quit 離開；/session 顯示 id。\n"
    );

    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() {
            continue;
        }
        if input == "/exit" || input == "/quit" {
            break;
        }
        if input == "/session" {
            println!("session {}", session.id);
            continue;
        }

        session.push(ChatMessage::user(input));
        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!(
            "[verdict: {} | session {} → {}]",
            verdict.kind(),
            session.id,
            path.display()
        );
        // Interactive: blocked does not exit the REPL — user can continue.
    }

    session.touch();
    let _ = save_session(&session)?;
    Ok(())
}

async fn cmd_smoke() -> Result<()> {
    println!("GrokBoy smoke (offline P6)…");
    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 + confirm + 6 browser incl. handoff)
    let defs = tool_definitions();
    let n_tools = defs.as_array().map(|a| a.len()).unwrap_or(0);
    if n_tools != 13 {
        return Err(anyhow!("expected 13 tool defs, got {n_tools}"));
    }
    let tool_names: Vec<&str> = defs
        .as_array()
        .unwrap()
        .iter()
        .filter_map(|t| t["function"]["name"].as_str())
        .collect();
    if !tool_names.contains(&"browser_handoff") {
        return Err(anyhow!("browser_handoff missing from tool defs"));
    }
    if !tool_names.contains(&"request_user_confirm") {
        return Err(anyhow!("request_user_confirm missing from tool defs"));
    }
    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");
    }

    // handoff protocol (auto-resume; no interactive stdin)
    // SAFETY: smoke is single-threaded for this env toggle.
    let prev_auto = std::env::var("GROKBOY_HANDOFF_AUTO").ok();
    unsafe { std::env::set_var("GROKBOY_HANDOFF_AUTO", "1") };
    let handoff = execute_tool(
        &ToolContext::new(repo.clone()),
        "browser_handoff",
        &json!({"reason": "smoke handoff check", "timeout_secs": 3}).to_string(),
    )
    .await;
    let handoff_v: serde_json::Value = serde_json::from_str(&handoff)?;
    // Fail-closed without Chromium, or resumed with snapshot if browsers installed.
    if handoff_v.get("error").is_none()
        && handoff_v.get("blocked").is_none()
        && handoff_v.get("handoff").and_then(|h| h.as_str()) != Some("resumed")
    {
        return Err(anyhow!("unexpected handoff smoke result: {handoff_v}"));
    }
    println!("  handoff auto   ok");
    // abort path
    unsafe { std::env::set_var("GROKBOY_HANDOFF_AUTO", "abort") };
    let abort_wait = grokboy_core::wait_for_handoff_resume(1);
    if !matches!(abort_wait, grokboy_core::HandoffWait::Aborted(_)) {
        return Err(anyhow!("expected HandoffWait::Aborted, got {abort_wait:?}"));
    }
    println!("  handoff abort  ok");
    match prev_auto {
        Some(v) => unsafe { std::env::set_var("GROKBOY_HANDOFF_AUTO", v) },
        None => unsafe { std::env::remove_var("GROKBOY_HANDOFF_AUTO") },
    }

    // request_user_confirm (auto approve/deny; no interactive stdin)
    let prev_confirm = std::env::var("GROKBOY_CONFIRM_AUTO").ok();
    unsafe { std::env::set_var("GROKBOY_CONFIRM_AUTO", "1") };
    let conf = execute_tool(
        &ctx,
        "request_user_confirm",
        &json!({
            "reason": "smoke confirm check",
            "prompt": "範例草稿",
            "timeout_secs": 2
        })
        .to_string(),
    )
    .await;
    let conf_v: serde_json::Value = serde_json::from_str(&conf)?;
    if conf_v["status"] != "approved" || conf_v["approved"] != true {
        return Err(anyhow!("request_user_confirm approve mismatch: {conf_v}"));
    }
    println!("  confirm auto   ok");
    unsafe { std::env::set_var("GROKBOY_CONFIRM_AUTO", "abort") };
    let deny = execute_tool(
        &ctx,
        "request_user_confirm",
        &json!({"reason": "smoke deny check", "timeout_secs": 2}).to_string(),
    )
    .await;
    let deny_v: serde_json::Value = serde_json::from_str(&deny)?;
    if deny_v["status"] != "denied" || deny_v["approved"] != false {
        return Err(anyhow!("request_user_confirm deny mismatch: {deny_v}"));
    }
    let deny_wait = grokboy_core::wait_for_user_confirm(1);
    if !matches!(deny_wait, grokboy_core::ConfirmWait::Denied(_)) {
        return Err(anyhow!("expected ConfirmWait::Denied, got {deny_wait:?}"));
    }
    println!("  confirm deny   ok");
    match prev_confirm {
        Some(v) => unsafe { std::env::set_var("GROKBOY_CONFIRM_AUTO", v) },
        None => unsafe { std::env::remove_var("GROKBOY_CONFIRM_AUTO") },
    }

    // 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");

    // agent CLI parse (offline)
    let (sid, help) = parse_agent_args(&[]).expect("empty agent args");
    if sid.is_some() || help {
        return Err(anyhow!("parse_agent_args([]) unexpected"));
    }
    let sess_args = vec!["--session".to_string(), "abc-123".to_string()];
    let (sid, help) = parse_agent_args(&sess_args)?;
    if sid.as_deref() != Some("abc-123") || help {
        return Err(anyhow!("parse_agent_args --session failed"));
    }
    let help_args = vec!["--help".to_string()];
    let (_, help) = parse_agent_args(&help_args)?;
    if !help {
        return Err(anyhow!("parse_agent_args --help failed"));
    }
    // one-turn session plumbing without API: system + user + fake assistant, save/load
    let mut agent_sess = Session::new(dir.clone());
    agent_sess.push(ChatMessage::system(AGENT_SYSTEM));
    agent_sess.push(ChatMessage::user("第一輪"));
    agent_sess.push(ChatMessage::assistant("回覆一"));
    agent_sess.push(ChatMessage::user("第二輪"));
    let agent_path = dir.join("agent-turn.json");
    let data = serde_json::to_vec_pretty(&agent_sess)?;
    std::fs::write(&agent_path, data)?;
    let loaded_agent: Session = serde_json::from_slice(&std::fs::read(&agent_path)?)?;
    if loaded_agent.messages.len() != 4 {
        return Err(anyhow!("agent multi-turn session plumbing failed"));
    }
    if loaded_agent.messages[0].text() != AGENT_SYSTEM {
        return Err(anyhow!("agent session missing system prompt"));
    }
    println!("  agent parse    ok");
    println!("  agent 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}")
}

#[cfg(test)]
mod tests {
    use super::parse_agent_args;

    #[test]
    fn agent_args_empty() {
        let (sid, help) = parse_agent_args(&[]).unwrap();
        assert!(sid.is_none());
        assert!(!help);
    }

    #[test]
    fn agent_args_session() {
        let args = vec!["--session".into(), "sess-1".into()];
        let (sid, help) = parse_agent_args(&args).unwrap();
        assert_eq!(sid.as_deref(), Some("sess-1"));
        assert!(!help);
    }

    #[test]
    fn agent_args_help() {
        let args = vec!["-h".into()];
        let (_, help) = parse_agent_args(&args).unwrap();
        assert!(help);
    }

    #[test]
    fn agent_args_rejects_extra() {
        let args = vec!["nope".into()];
        assert!(parse_agent_args(&args).is_err());
    }
}