LazyBoy2/crates/grokboy-core/src/agent.rs

//! Multi-step OpenAI-compatible tool-calling ReAct loop (P2: completion, loop guard, truncation).

use crate::config::Config;
use crate::model::{ChatMessage, Role, ToolCall, chat_completion};
use crate::tools::{ToolContext, execute_tool, is_completion_tool, tool_definitions};
use anyhow::Result;
use serde_json::Value;
use std::future::Future;

pub const DEFAULT_MAX_ROUNDS: usize = 12;
pub const DEFAULT_CONTEXT_CHARS: usize = 100_000;
pub const LOOP_GUARD_REPEAT: usize = 3;

pub const AGENT_SYSTEM: &str = "\
You are GrokBoy, a concise local coding assistant with tools.
Use tools when they help solve the task; otherwise answer directly.
Prefer short, clear answers. Traditional Chinese is welcome when the user writes in Chinese.
Available tools: shell, list_dir, read_file, write_file, report_done, report_blocked, and optional browser_* (Playwright DOM: browser_navigate, browser_snapshot, browser_click, browser_type, browser_eval).
For web pages prefer DOM snapshot + selector/role click/type — not screenshots or pixel XY clicks.
When the task is finished, call report_done with a short summary.
If you are stuck or cannot proceed, call report_blocked with the reason — do not invent results or loop.
Do not invent tool results — call the tools.";

/// Final verdict from the agent loop (fail-closed when stuck).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AgentVerdict {
    /// Model called `report_done`.
    Done(String),
    /// Model called `report_blocked`, loop guard, or max rounds.
    Blocked(String),
    /// Model returned final text without a completion tool.
    Answer(String),
}

impl AgentVerdict {
    pub fn message(&self) -> &str {
        match self {
            Self::Done(s) | Self::Blocked(s) | Self::Answer(s) => s,
        }
    }

    pub fn kind(&self) -> &'static str {
        match self {
            Self::Done(_) => "done",
            Self::Blocked(_) => "blocked",
            Self::Answer(_) => "answer",
        }
    }
}

/// Resolve context budget from `GROKBOY_CONTEXT_CHARS` or default.
pub fn context_char_budget() -> usize {
    std::env::var("GROKBOY_CONTEXT_CHARS")
        .ok()
        .and_then(|s| s.parse().ok())
        .filter(|&n| n > 0)
        .unwrap_or(DEFAULT_CONTEXT_CHARS)
}

/// Stable signature for a single tool call (name + args).
pub fn tool_call_signature(call: &ToolCall) -> String {
    format!("{}:{}", call.function.name, call.function.arguments)
}

/// Signature for a whole round of tool calls (order-preserving).
pub fn round_signature(calls: &[ToolCall]) -> String {
    calls
        .iter()
        .map(tool_call_signature)
        .collect::<Vec<_>>()
        .join("\n")
}

/// Approximate serialized size of one message (chars).
pub fn message_char_len(msg: &ChatMessage) -> usize {
    let mut n = 8; // role overhead
    if let Some(c) = &msg.content {
        n += c.len();
    }
    if let Some(id) = &msg.tool_call_id {
        n += id.len();
    }
    if let Some(calls) = &msg.tool_calls {
        for c in calls {
            n += c.id.len() + c.function.name.len() + c.function.arguments.len() + 16;
        }
    }
    n
}

pub fn messages_char_len(messages: &[ChatMessage]) -> usize {
    messages.iter().map(message_char_len).sum()
}

/// Truncate older tool results / middle turns while keeping system + recent tail + last user goal.
/// Mutates in place. Keeps it simple: shrink old tool contents first, then drop middle messages.
pub fn truncate_messages(messages: &mut Vec<ChatMessage>, budget: usize) {
    if budget == 0 || messages_char_len(messages) <= budget {
        return;
    }

    // 1) Shrink older tool message contents (keep last few tool results intact).
    let tool_indices: Vec<usize> = messages
        .iter()
        .enumerate()
        .filter(|(_, m)| m.role == Role::Tool)
        .map(|(i, _)| i)
        .collect();
    let keep_recent_tools = 4usize;
    let shrink_until = tool_indices.len().saturating_sub(keep_recent_tools);
    for &idx in tool_indices.iter().take(shrink_until) {
        if let Some(content) = messages[idx].content.as_mut() {
            if content.len() > 120 {
                let omitted = content.len();
                *content = format!("[truncated tool result; was {omitted} chars]");
            }
        }
        if messages_char_len(messages) <= budget {
            return;
        }
    }

    // 2) Drop middle messages while preserving: leading system*, last user, and a recent tail.
    while messages_char_len(messages) > budget && messages.len() > 4 {
        let first_drop = messages
            .iter()
            .enumerate()
            .find(|(i, m)| *i > 0 && m.role != Role::System)
            .map(|(i, _)| i);
        let Some(i) = first_drop else { break };

        // Never drop the last user message or the last two messages.
        let last_user = messages
            .iter()
            .rposition(|m| m.role == Role::User)
            .unwrap_or(messages.len());
        if i >= last_user || i + 2 >= messages.len() {
            // Shrink remaining large contents instead.
            for msg in messages.iter_mut() {
                if let Some(content) = msg.content.as_mut() {
                    if content.len() > 200 {
                        let keep = 200.min(content.len());
                        let omitted = content.len().saturating_sub(keep);
                        *content = format!("{}…\n[truncated {omitted} chars]", &content[..keep]);
                    }
                }
            }
            break;
        }

        // If dropping an assistant with tool_calls, also drop following tool messages for those ids.
        let drop_ids: Vec<String> = messages
            .get(i)
            .and_then(|m| m.tool_calls.as_ref())
            .map(|calls| calls.iter().map(|c| c.id.clone()).collect())
            .unwrap_or_default();

        messages.remove(i);
        let j = i;
        while j < messages.len() {
            let is_orphan_tool = messages[j].role == Role::Tool
                && messages[j]
                    .tool_call_id
                    .as_ref()
                    .is_some_and(|id| drop_ids.contains(id));
            if is_orphan_tool {
                messages.remove(j);
            } else {
                break;
            }
        }
    }
}

/// Run the agent with a live HTTP model client.
pub async fn run_agent(
    config: &Config,
    messages: &mut Vec<ChatMessage>,
    tool_ctx: &ToolContext,
    max_rounds: usize,
) -> Result<AgentVerdict> {
    let budget = context_char_budget();
    let config = config.clone();
    run_agent_with(messages, tool_ctx, max_rounds, budget, move |msgs, tools| {
        let config = config.clone();
        async move { chat_completion(&config, &msgs, Some(&tools)).await }
    })
    .await
}

/// Core loop with injectable completer (for offline tests).
/// `complete` receives a snapshot of messages and tool defs each round.
pub async fn run_agent_with<F, Fut>(
    messages: &mut Vec<ChatMessage>,
    tool_ctx: &ToolContext,
    max_rounds: usize,
    context_budget: usize,
    mut complete: F,
) -> Result<AgentVerdict>
where
    F: FnMut(Vec<ChatMessage>, Value) -> Fut,
    Fut: Future<Output = Result<ChatMessage>>,
{
    let tools = tool_definitions();
    let mut prev_round_sig: Option<String> = None;
    let mut same_sig_streak: usize = 0;

    for round in 0..max_rounds {
        truncate_messages(messages, context_budget);

        let reply = complete(messages.clone(), tools.clone()).await?;
        let tool_calls = reply.tool_calls.clone().unwrap_or_default();

        if tool_calls.is_empty() {
            let last_text = reply.text().to_string();
            messages.push(reply);
            if last_text.trim().is_empty() {
                return Ok(AgentVerdict::Blocked(
                    "model returned empty final answer".into(),
                ));
            }
            return Ok(AgentVerdict::Answer(last_text));
        }

        // Loop guard: identical tool-call round repeated N times → fail closed.
        let sig = round_signature(&tool_calls);
        if prev_round_sig.as_deref() == Some(sig.as_str()) {
            same_sig_streak += 1;
        } else {
            same_sig_streak = 1;
            prev_round_sig = Some(sig);
        }
        if same_sig_streak >= LOOP_GUARD_REPEAT {
            messages.push(reply);
            let reason = format!(
                "loop guard: identical tool calls repeated {LOOP_GUARD_REPEAT} times (round {})",
                round + 1
            );
            return Ok(AgentVerdict::Blocked(reason));
        }

        messages.push(reply);

        let mut completion: Option<AgentVerdict> = None;
        for call in &tool_calls {
            let result =
                execute_tool(tool_ctx, &call.function.name, &call.function.arguments).await;
            if completion.is_none() && is_completion_tool(&call.function.name) {
                completion = parse_completion_verdict(&call.function.name, &result);
            }
            messages.push(ChatMessage::tool(&call.id, result));
        }

        if let Some(verdict) = completion {
            return Ok(verdict);
        }
    }

    Ok(AgentVerdict::Blocked(format!(
        "blocked: reached max rounds ({max_rounds}) without completion"
    )))
}

fn parse_completion_verdict(name: &str, result_json: &str) -> Option<AgentVerdict> {
    let v: Value = serde_json::from_str(result_json).ok()?;
    if v.get("error").is_some() {
        return Some(AgentVerdict::Blocked(format!(
            "completion tool {name} failed: {}",
            v["error"]
        )));
    }
    match name {
        "report_done" => {
            let msg = v
                .get("message")
                .and_then(|x| x.as_str())
                .unwrap_or("done")
                .to_string();
            Some(AgentVerdict::Done(msg))
        }
        "report_blocked" => {
            let reason = v
                .get("reason")
                .and_then(|x| x.as_str())
                .unwrap_or("blocked")
                .to_string();
            Some(AgentVerdict::Blocked(reason))
        }
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::{FunctionCall, ToolCall};
    use serde_json::json;
    use std::sync::{Arc, Mutex};

    fn tc(id: &str, name: &str, args: &str) -> ToolCall {
        ToolCall {
            id: id.into(),
            kind: "function".into(),
            function: FunctionCall {
                name: name.into(),
                arguments: args.into(),
            },
        }
    }

    #[test]
    fn system_prompt_mentions_tools_and_chinese() {
        assert!(AGENT_SYSTEM.contains("shell"));
        assert!(AGENT_SYSTEM.contains("report_done"));
        assert!(AGENT_SYSTEM.contains("report_blocked"));
        assert!(AGENT_SYSTEM.contains("browser_navigate") || AGENT_SYSTEM.contains("browser_"));
        assert!(AGENT_SYSTEM.contains("Traditional Chinese") || AGENT_SYSTEM.contains("Chinese"));
    }

    #[test]
    fn default_max_rounds_is_12() {
        assert_eq!(DEFAULT_MAX_ROUNDS, 12);
    }

    #[test]
    fn round_signature_stable() {
        let a = vec![tc("1", "list_dir", r#"{"path":"."}"#)];
        let b = vec![tc("2", "list_dir", r#"{"path":"."}"#)];
        assert_eq!(round_signature(&a), round_signature(&b));
        let c = vec![tc("1", "list_dir", r#"{"path":"src"}"#)];
        assert_ne!(round_signature(&a), round_signature(&c));
    }

    #[test]
    fn truncate_keeps_system_and_shrinks_old_tools() {
        let mut msgs = vec![
            ChatMessage::system("sys"),
            ChatMessage::user("goal"),
            ChatMessage::assistant_tool_calls(None, vec![tc("c1", "read_file", r#"{"path":"a"}"#)]),
            ChatMessage::tool("c1", "X".repeat(5_000)),
            ChatMessage::assistant_tool_calls(None, vec![tc("c2", "read_file", r#"{"path":"b"}"#)]),
            ChatMessage::tool("c2", "Y".repeat(5_000)),
            ChatMessage::assistant_tool_calls(None, vec![tc("c3", "read_file", r#"{"path":"c"}"#)]),
            ChatMessage::tool("c3", "Z".repeat(5_000)),
            ChatMessage::assistant_tool_calls(None, vec![tc("c4", "read_file", r#"{"path":"d"}"#)]),
            ChatMessage::tool("c4", "W".repeat(5_000)),
            ChatMessage::assistant_tool_calls(None, vec![tc("c5", "read_file", r#"{"path":"e"}"#)]),
            ChatMessage::tool("c5", "V".repeat(5_000)),
            ChatMessage::user("still the goal"),
            ChatMessage::assistant("recent"),
        ];
        let before = messages_char_len(&msgs);
        truncate_messages(&mut msgs, 8_000);
        let after = messages_char_len(&msgs);
        assert!(after < before);
        assert_eq!(msgs[0].text(), "sys");
        assert!(msgs.iter().any(|m| m.role == Role::User));
        let old_tool = msgs
            .iter()
            .find(|m| m.role == Role::Tool && m.tool_call_id.as_deref() == Some("c1"));
        if let Some(t) = old_tool {
            assert!(
                t.text().contains("truncated") || t.text().len() < 5_000,
                "{}",
                t.text()
            );
        }
    }

    #[tokio::test]
    async fn loop_stops_on_report_done() {
        let dir = std::env::temp_dir().join(format!("grokboy-agent-done-{}", std::process::id()));
        let _ = std::fs::create_dir_all(&dir);
        let ctx = ToolContext::new(dir.clone());
        let mut messages = vec![
            ChatMessage::system(AGENT_SYSTEM),
            ChatMessage::user("finish"),
        ];
        let calls = Arc::new(Mutex::new(0usize));
        let calls2 = calls.clone();

        let verdict = run_agent_with(&mut messages, &ctx, 5, 100_000, move |_msgs, _tools| {
            let n = {
                let mut g = calls2.lock().unwrap();
                *g += 1;
                *g
            };
            async move {
                if n == 1 {
                    Ok(ChatMessage::assistant_tool_calls(
                        None,
                        vec![tc(
                            "done1",
                            "report_done",
                            &json!({"message": "任務完成"}).to_string(),
                        )],
                    ))
                } else {
                    Ok(ChatMessage::assistant("should not reach"))
                }
            }
        })
        .await
        .unwrap();

        assert_eq!(verdict, AgentVerdict::Done("任務完成".into()));
        assert_eq!(*calls.lock().unwrap(), 1);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn loop_stops_on_report_blocked() {
        let dir = std::env::temp_dir().join(format!("grokboy-agent-blk-{}", std::process::id()));
        let _ = std::fs::create_dir_all(&dir);
        let ctx = ToolContext::new(dir.clone());
        let mut messages = vec![ChatMessage::user("x")];

        let verdict = run_agent_with(&mut messages, &ctx, 5, 100_000, move |_msgs, _tools| async move {
            Ok(ChatMessage::assistant_tool_calls(
                None,
                vec![tc(
                    "b1",
                    "report_blocked",
                    &json!({"reason": "permission denied"}).to_string(),
                )],
            ))
        })
        .await
        .unwrap();

        assert_eq!(verdict, AgentVerdict::Blocked("permission denied".into()));
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn loop_guard_blocks_repeated_calls() {
        let dir = std::env::temp_dir().join(format!("grokboy-agent-loop-{}", std::process::id()));
        let _ = std::fs::create_dir_all(&dir);
        let ctx = ToolContext::new(dir.clone());
        std::fs::write(dir.join("f.txt"), "hi").unwrap();
        let mut messages = vec![ChatMessage::user("loop")];
        let args = json!({"path": "f.txt"}).to_string();

        let verdict = run_agent_with(&mut messages, &ctx, 12, 100_000, move |_msgs, _tools| {
            let args = args.clone();
            async move {
                Ok(ChatMessage::assistant_tool_calls(
                    None,
                    vec![tc("r1", "read_file", &args)],
                ))
            }
        })
        .await
        .unwrap();

        match verdict {
            AgentVerdict::Blocked(reason) => {
                assert!(reason.contains("loop guard"), "{reason}");
            }
            other => panic!("expected blocked, got {other:?}"),
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn max_rounds_fail_closed() {
        let dir = std::env::temp_dir().join(format!("grokboy-agent-max-{}", std::process::id()));
        let _ = std::fs::create_dir_all(&dir);
        let ctx = ToolContext::new(dir.clone());
        std::fs::write(dir.join("a.txt"), "1").unwrap();
        std::fs::write(dir.join("b.txt"), "2").unwrap();
        let mut messages = vec![ChatMessage::user("spin")];
        let n = Arc::new(Mutex::new(0usize));
        let n2 = n.clone();

        let verdict = run_agent_with(&mut messages, &ctx, 2, 100_000, move |_msgs, _tools| {
            let i = {
                let mut g = n2.lock().unwrap();
                *g += 1;
                *g
            };
            async move {
                let path = if i == 1 { "a.txt" } else { "b.txt" };
                Ok(ChatMessage::assistant_tool_calls(
                    None,
                    vec![tc(
                        &format!("c{i}"),
                        "read_file",
                        &json!({"path": path}).to_string(),
                    )],
                ))
            }
        })
        .await
        .unwrap();

        match verdict {
            AgentVerdict::Blocked(reason) => {
                assert!(reason.contains("max rounds"), "{reason}");
            }
            other => panic!("expected blocked, got {other:?}"),
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn plain_answer_without_tools() {
        let dir = std::env::temp_dir().join(format!("grokboy-agent-ans-{}", std::process::id()));
        let _ = std::fs::create_dir_all(&dir);
        let ctx = ToolContext::new(dir.clone());
        let mut messages = vec![ChatMessage::user("hi")];

        let verdict = run_agent_with(&mut messages, &ctx, 5, 100_000, move |_msgs, _tools| async move {
            Ok(ChatMessage::assistant("hello there"))
        })
        .await
        .unwrap();

        assert_eq!(verdict, AgentVerdict::Answer("hello there".into()));
        let _ = std::fs::remove_dir_all(&dir);
    }
}