wasm-agent 0.2.0

// SPDX-License-Identifier: MIT
//! ReAct loop runner: Thought -> Action -> Observation -> FinalAnswer.

use crate::error::AgentError;
use crate::history::ConversationHistory;
use crate::tools::ToolRegistry;
use crate::types::{AgentConfig, Message, ReActStep, Role};
use serde::{Deserialize, Serialize};

/// Find `label` at the start of a line (case-insensitive, leading spaces and
/// `*`/`#` markdown decoration allowed). Returns the byte offset just after the
/// label's colon.
fn find_label(text: &str, label: &str) -> Option<(usize, usize)> {
    let mut offset = 0;
    for line in text.split_inclusive('\n') {
        let trimmed = line.trim_start_matches([' ', '\t', '*', '#']);
        let lead = line.len() - trimmed.len();
        if trimmed.len() >= label.len()
            && trimmed.is_char_boundary(label.len())
            && trimmed[..label.len()].eq_ignore_ascii_case(label)
        {
            let after = trimmed[label.len()..].trim_start_matches('*');
            if let Some(rest) = after.strip_prefix(':') {
                let rest = rest.trim_start_matches('*');
                let start = offset + lead;
                let content = offset + line.len() - rest.len();
                return Some((start, content));
            }
        }
        offset += line.len();
    }
    None
}

/// The text after a label up to the next ReAct label (or the end).
fn section<'a>(text: &'a str, content_start: usize, stop_labels: &[&str]) -> &'a str {
    let rest = &text[content_start..];
    let end = stop_labels
        .iter()
        .filter_map(|l| find_label(rest, l).map(|(s, _)| s))
        .min()
        .unwrap_or(rest.len());
    rest[..end].trim()
}

const LABELS: &[&str] = &["Thought", "Action Input", "Action", "Observation", "Final Answer"];

/// Parses a raw LLM response into a [`ReActStep`].
///
/// Models usually write several labelled lines in one response, for example
/// `Thought: ...` followed by `Action: ...`. The most decisive one wins:
/// `Final Answer` over `Action` over `Thought`. Labels must start a line and
/// are matched case-insensitively (`**Action:**` also works). Anything from an
/// `Observation:` line on is ignored, because the model is inventing a tool
/// result it has not received.
///
/// Supported action forms:
/// - `Action: tool(input)` (one trailing `)` is removed, inner parentheses are kept)
/// - `Action: tool: input`
/// - `Action: tool` followed by an `Action Input: input` line (LangChain style)
/// - `Action: tool` with no input
///
/// # Errors
/// Returns [`AgentError::ParseError`] if no label is found, or `Action:` names no tool.
pub fn parse_react_step(response: &str) -> Result<ReActStep, AgentError> {
    // Drop a hallucinated observation and everything after it.
    let text = match find_label(response, "Observation") {
        Some((start, _)) => &response[..start],
        None => response,
    };

    if let Some((_, c)) = find_label(text, "Final Answer") {
        return Ok(ReActStep::FinalAnswer(text[c..].trim().to_string()));
    }

    if let Some((_, c)) = find_label(text, "Action") {
        // "Action Input:" also starts with "Action"; find_label requires the
        // colon right after the label, so it never matches here.
        let rest = section(text, c, LABELS);
        let explicit_input = find_label(text, "Action Input").map(|(_, ic)| section(text, ic, LABELS));
        let (tool, input) = split_action(rest);
        if tool.is_empty() {
            return Err(AgentError::ParseError(format!("Action without a tool name: {}", response.trim())));
        }
        let input = match explicit_input {
            Some(i) if input.is_empty() => i.to_string(),
            _ => input,
        };
        return Ok(ReActStep::Action { tool, input });
    }

    if let Some((_, c)) = find_label(text, "Thought") {
        return Ok(ReActStep::Thought(section(text, c, LABELS).to_string()));
    }

    Err(AgentError::ParseError(format!("Could not parse ReAct step from: {}", response.trim())))
}

fn split_action(rest: &str) -> (String, String) {
    let first_line = rest.lines().next().unwrap_or("").trim();
    let paren = first_line.find('(');
    let colon = first_line.find(':');
    match (paren, colon) {
        (Some(p), c) if c.is_none_or(|c| p < c) => {
            let tool = first_line[..p].trim().to_string();
            // Input may span lines: take everything after "(" in the section.
            let after = rest[rest.find('(').map_or(0, |i| i + 1)..].trim_end();
            let input = after.strip_suffix(')').unwrap_or(after).trim().to_string();
            (tool, input)
        }
        (_, Some(c)) => {
            let tool = first_line[..c].trim().to_string();
            let after_colon = &rest[rest.find(':').map_or(rest.len(), |i| i + 1)..];
            (tool, after_colon.trim().to_string())
        }
        _ => (first_line.trim_matches('`').to_string(), rest[first_line.len().min(rest.len())..].trim().to_string()),
    }
}

/// Loop state shared by [`LoopRunner`] (borrows its config and tools) and
/// [`Agent`] (owns them).
#[derive(Debug)]
struct State {
    history: ConversationHistory,
    steps: u32,
    started_at_ms: Option<f64>,
}

impl State {
    fn new(config: &AgentConfig) -> Self {
        Self { history: ConversationHistory::new(config.context_token_limit), steps: 0, started_at_ms: None }
    }

    fn check_budget(&mut self, config: &AgentConfig, now_ms: Option<f64>) -> Result<(), AgentError> {
        if self.steps >= config.max_iterations {
            return Err(AgentError::MaxIterationsExceeded(config.max_iterations));
        }
        if let (Some(now), Some(limit)) = (now_ms, config.timeout_ms) {
            let start = *self.started_at_ms.get_or_insert(now);
            if now - start > limit as f64 {
                return Err(AgentError::Timeout { timeout_ms: limit });
            }
        }
        Ok(())
    }

    fn observation(result: crate::types::ToolResult) -> String {
        if result.success { result.output } else { format!("Error: {}", result.output) }
    }

    fn step(&mut self, config: &AgentConfig, registry: &ToolRegistry, llm_response: &str, now_ms: Option<f64>) -> Result<ReActStep, AgentError> {
        self.check_budget(config, now_ms)?;
        let parsed = parse_react_step(llm_response);
        let (step, obs) = match parsed {
            Ok(step) => {
                let obs = match &step {
                    ReActStep::Action { tool, input } => match registry.dispatch(tool, input) {
                        Ok(result) => Some(Self::observation(result)),
                        Err(AgentError::ToolNotFound { name }) if config.recover_from_errors => Some(format!(
                            "Error: there is no tool named '{name}'. Available tools: {}.",
                            registry.tool_names().join(", ")
                        )),
                        Err(e) => return Err(e),
                    },
                    _ => None,
                };
                (step, obs)
            }
            Err(AgentError::ParseError(_)) if config.recover_from_errors => (
                ReActStep::Observation("format error".into()),
                Some("Error: reply with exactly one of 'Thought: ...', 'Action: tool(input)' or 'Final Answer: ...'.".into()),
            ),
            Err(e) => return Err(e),
        };
        self.steps += 1;
        // Keep only the part the parser used: a hallucinated "Observation:"
        // must not end up in the history as if it were real.
        let kept = match find_label(llm_response, "Observation") {
            Some((start, _)) => llm_response[..start].trim_end(),
            None => llm_response,
        };
        self.history.push_with_eviction(Message::assistant(kept));
        if let Some(obs) = obs {
            self.history.push_with_eviction(Message::new(Role::Tool, obs));
        }
        Ok(step)
    }

    fn step_response(
        &mut self,
        config: &AgentConfig,
        registry: &ToolRegistry,
        response_json: &str,
        now_ms: Option<f64>,
    ) -> Result<NativeStep, AgentError> {
        self.check_budget(config, now_ms)?;
        let turn = crate::native::parse_model_response(response_json)?;
        self.steps += 1;
        let text = turn.text.clone().unwrap_or_default();
        if turn.calls.is_empty() {
            self.history.push_with_eviction(Message::assistant(text.clone()));
            return Ok(NativeStep::Final(text));
        }
        let mut assistant = Message::assistant(text);
        assistant.tool_calls = turn.calls.clone();
        self.history.push_with_eviction(assistant);
        let mut outputs = Vec::with_capacity(turn.calls.len());
        for call in &turn.calls {
            let output = match registry.dispatch_call(call) {
                Ok(result) => crate::native::ToolOutput { id: call.id.clone(), name: call.name.clone(), success: result.success, output: result.output },
                Err(AgentError::ToolNotFound { name }) if config.recover_from_errors => crate::native::ToolOutput {
                    id: call.id.clone(),
                    name: call.name.clone(),
                    success: false,
                    output: format!("there is no tool named '{name}'"),
                },
                Err(e) => return Err(e),
            };
            let mut m = Message::new(Role::Tool, Self::observation(crate::types::ToolResult {
                tool_name: output.name.clone(),
                output: output.output.clone(),
                success: output.success,
            }));
            m.tool_call_id = output.id.clone();
            self.history.push_with_eviction(m);
            outputs.push(output);
        }
        Ok(NativeStep::ToolCalls(outputs))
    }

    fn transcript(&self) -> String {
        let mut out = String::new();
        for m in self.history.messages() {
            match m.role {
                Role::Tool => {
                    out.push_str("Observation: ");
                    out.push_str(&m.content);
                }
                Role::User => {
                    out.push_str("Question: ");
                    out.push_str(&m.content);
                }
                _ => out.push_str(&m.content),
            }
            out.push('\n');
        }
        out
    }
}

/// What one native tool-calling turn did (see [`LoopRunner::step_response`]).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum NativeStep {
    /// The model called tools; here are their results, in call order. Send
    /// them back to the model (see [`crate::native::tool_results_openai`] and
    /// [`crate::native::tool_results_anthropic`]) and step again.
    ToolCalls(Vec<crate::native::ToolOutput>),
    /// The model answered without calling a tool.
    Final(String),
}

/// The ReAct loop execution context.
///
/// Holds references to the agent configuration and tool registry and drives
/// the Thought -> Action -> Observation -> FinalAnswer cycle. It never calls a
/// model itself: you pass each model response to [`step`](Self::step) (text
/// ReAct) or [`step_response`](Self::step_response) (native tool calling).
/// Use [`Agent`] if you would rather not manage the borrows.
pub struct LoopRunner<'a> {
    config: &'a AgentConfig,
    registry: &'a ToolRegistry,
    state: State,
}

impl<'a> LoopRunner<'a> {
    /// Creates a new runner with its own conversation history.
    pub fn new(config: &'a AgentConfig, registry: &'a ToolRegistry) -> Self {
        Self { config, registry, state: State::new(config) }
    }

    /// Adds the standard ReAct system prompt (format rules plus the tool list)
    /// as the first message. System messages are never evicted from history.
    pub fn with_system_prompt(mut self) -> Self {
        self.state.history.push_with_eviction(Message::system(system_prompt(self.registry)));
        self
    }

    /// Adds the user's question to the history.
    pub fn ask(&mut self, question: &str) {
        self.state.history.push_with_eviction(Message::user(question));
    }

    /// Processes one text LLM response, dispatching the tool when it is an Action.
    ///
    /// With `config.recover_from_errors`, an unknown tool or an unparseable
    /// reply becomes an `Error: ...` observation the model can correct,
    /// instead of an error that ends the loop.
    ///
    /// # Errors
    /// - [`AgentError::MaxIterationsExceeded`]: more than `config.max_iterations` steps.
    /// - [`AgentError::ParseError`]: the response does not match any ReAct format.
    /// - [`AgentError::ToolNotFound`]: the requested tool is not in the registry.
    pub fn step(&mut self, llm_response: &str) -> Result<ReActStep, AgentError> {
        self.state.step(self.config, self.registry, llm_response, None)
    }

    /// Like [`step`](Self::step), and also enforces `config.timeout_ms`. Pass
    /// the current time in milliseconds (for example `Date.now()`); the first
    /// call starts the clock.
    ///
    /// # Errors
    /// [`AgentError::Timeout`] once more than `timeout_ms` has passed since the
    /// first call, plus everything [`step`](Self::step) returns.
    pub fn step_at(&mut self, llm_response: &str, now_ms: f64) -> Result<ReActStep, AgentError> {
        self.state.step(self.config, self.registry, llm_response, Some(now_ms))
    }

    /// Processes one native tool-calling response: the raw JSON body from
    /// OpenAI Chat Completions or Responses, Anthropic Messages, or Ollama's
    /// chat API. Every tool call in it is dispatched (arguments arrive as JSON;
    /// see [`ToolRegistry::dispatch_call`]) and the results are returned and
    /// kept in the history with their call ids.
    ///
    /// # Errors
    /// [`AgentError::ParseError`] if the body is not a recognised response,
    /// [`AgentError::MaxIterationsExceeded`], [`AgentError::ToolNotFound`]
    /// (unless `recover_from_errors`).
    pub fn step_response(&mut self, response_json: &str) -> Result<NativeStep, AgentError> {
        self.state.step_response(self.config, self.registry, response_json, None)
    }

    /// Runs a pre-scripted sequence of LLM responses, returning the final answer.
    ///
    /// Useful for offline simulation and deterministic testing.
    ///
    /// # Errors
    /// - [`AgentError::MaxIterationsExceeded`]: sequence exhausted before a `FinalAnswer`.
    /// - Any error from [`step`](Self::step).
    pub fn run_scripted(&mut self, responses: &[&str]) -> Result<String, AgentError> {
        for response in responses {
            if let ReActStep::FinalAnswer(answer) = self.step(response)? {
                return Ok(answer);
            }
        }
        Err(AgentError::MaxIterationsExceeded(self.config.max_iterations))
    }

    /// The conversation so far as one prompt string, ready to send to a
    /// completion-style model: system and user messages as-is, model turns
    /// as-is, tool results as `Observation: ...` lines.
    pub fn transcript(&self) -> String {
        self.state.transcript()
    }

    /// Returns a reference to the accumulated conversation history.
    pub fn history(&self) -> &ConversationHistory { &self.state.history }

    /// Number of model responses processed so far.
    pub fn iteration_count(&self) -> usize { self.state.steps as usize }
}

/// A ReAct agent that owns its configuration, tools and history (no
/// lifetimes to manage). Same behaviour as [`LoopRunner`].
///
/// ```
/// use wasm_agent::{Agent, AgentConfig, ReActStep, ToolSpec};
/// use wasm_agent::types::ToolResult;
///
/// let mut agent = Agent::new(AgentConfig::default());
/// agent.tools_mut().register(
///     ToolSpec::new("echo", "Repeats its input", r#"{"type":"string"}"#),
///     Box::new(|s: &str| ToolResult { tool_name: "echo".into(), output: s.into(), success: true }),
/// )?;
/// agent.step("Action: echo(hi)")?;
/// assert_eq!(agent.step("Final Answer: done")?, ReActStep::FinalAnswer("done".into()));
/// # Ok::<(), wasm_agent::AgentError>(())
/// ```
pub struct Agent {
    config: AgentConfig,
    registry: ToolRegistry,
    state: State,
}

impl Agent {
    /// An agent with no tools yet.
    pub fn new(config: AgentConfig) -> Self {
        let state = State::new(&config);
        Self { config, registry: ToolRegistry::new(), state }
    }

    /// An agent with the given tools.
    pub fn with_tools(config: AgentConfig, registry: ToolRegistry) -> Self {
        let state = State::new(&config);
        Self { config, registry, state }
    }

    /// The tool registry, to register more tools.
    pub fn tools_mut(&mut self) -> &mut ToolRegistry { &mut self.registry }

    /// The tool registry.
    pub fn tools(&self) -> &ToolRegistry { &self.registry }

    /// The configuration.
    pub fn config(&self) -> &AgentConfig { &self.config }

    /// Adds the standard ReAct system prompt as the first message.
    pub fn add_system_prompt(&mut self) {
        let prompt = system_prompt(&self.registry);
        self.state.history.push_with_eviction(Message::system(prompt));
    }

    /// Adds the user's question to the history.
    pub fn ask(&mut self, question: &str) {
        self.state.history.push_with_eviction(Message::user(question));
    }

    /// See [`LoopRunner::step`].
    ///
    /// # Errors
    /// See [`LoopRunner::step`].
    pub fn step(&mut self, llm_response: &str) -> Result<ReActStep, AgentError> {
        self.state.step(&self.config, &self.registry, llm_response, None)
    }

    /// See [`LoopRunner::step_at`].
    ///
    /// # Errors
    /// See [`LoopRunner::step_at`].
    pub fn step_at(&mut self, llm_response: &str, now_ms: f64) -> Result<ReActStep, AgentError> {
        self.state.step(&self.config, &self.registry, llm_response, Some(now_ms))
    }

    /// See [`LoopRunner::step_response`].
    ///
    /// # Errors
    /// See [`LoopRunner::step_response`].
    pub fn step_response(&mut self, response_json: &str) -> Result<NativeStep, AgentError> {
        self.state.step_response(&self.config, &self.registry, response_json, None)
    }

    /// See [`LoopRunner::transcript`].
    pub fn transcript(&self) -> String { self.state.transcript() }

    /// The conversation history.
    pub fn history(&self) -> &ConversationHistory { &self.state.history }

    /// Number of model responses processed so far.
    pub fn iteration_count(&self) -> usize { self.state.steps as usize }
}

/// The standard ReAct instructions plus the registry's tool list, for use as
/// a system prompt.
pub fn system_prompt(registry: &ToolRegistry) -> String {
    format!(
        "Answer the question. You can use tools.\n\n{}\n\n\
         Use this format, one step per reply:\n\
         Thought: what you are thinking\n\
         Action: tool_name(input)\n\
         Then wait: the tool result comes back as \"Observation: ...\".\n\
         When you know the answer, reply with:\n\
         Final Answer: the answer",
        registry.tools_prompt()
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tools::{ToolRegistry, ToolSpec};
    use crate::types::ToolResult;

    fn registry_with_echo() -> ToolRegistry {
        let mut reg = ToolRegistry::new();
        reg.register(
            ToolSpec::new("echo", "Echoes input", "{}"),
            Box::new(|input: &str| ToolResult {
                tool_name: "echo".into(),
                output: format!("echoed: {input}"),
                success: true,
            }),
        ).unwrap();
        reg
    }

    #[test]
    fn test_parse_thought() {
        let step = parse_react_step("Thought: I need to search for information").unwrap();
        assert!(matches!(step, ReActStep::Thought(_)));
    }

    #[test]
    fn test_parse_action_paren_syntax() {
        let step = parse_react_step("Action: echo(hello world)").unwrap();
        assert!(matches!(&step, ReActStep::Action { tool, input }
            if tool == "echo" && input == "hello world"));
    }

    #[test]
    fn test_parse_action_colon_syntax() {
        let step = parse_react_step("Action: echo: hello").unwrap();
        assert!(matches!(&step, ReActStep::Action { tool, .. } if tool == "echo"));
    }

    #[test]
    fn test_parse_final_answer() {
        let step = parse_react_step("Final Answer: 42").unwrap();
        assert!(matches!(&step, ReActStep::FinalAnswer(s) if s == "42"));
    }

    #[test]
    fn test_parse_unknown_format_returns_parse_error() {
        let err = parse_react_step("This is just random text without a prefix").unwrap_err();
        assert!(matches!(err, AgentError::ParseError(_)));
    }

    #[test]
    fn test_parse_action_no_args() {
        let step = parse_react_step("Action: mytool").unwrap();
        assert!(matches!(&step, ReActStep::Action { tool, .. } if tool == "mytool"));
    }

    #[test]
    fn test_loop_runner_scripted_final_answer() {
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        let answer = runner.run_scripted(&[
            "Thought: Let me think",
            "Action: echo(test)",
            "Final Answer: done",
        ]).unwrap();
        assert_eq!(answer, "done");
    }

    #[test]
    fn test_loop_runner_max_iterations_returns_error() {
        let config = AgentConfig { max_iterations: 2, ..Default::default() };
        let reg = ToolRegistry::new();
        let mut runner = LoopRunner::new(&config, &reg);
        let err = runner.run_scripted(&[
            "Thought: step 1",
            "Thought: step 2",
            "Thought: step 3",
        ]).unwrap_err();
        assert!(matches!(err, AgentError::MaxIterationsExceeded(2)));
    }

    #[test]
    fn test_loop_runner_tool_not_found_returns_error() {
        let config = AgentConfig::default();
        let reg = ToolRegistry::new();
        let mut runner = LoopRunner::new(&config, &reg);
        let err = runner.step("Action: missing_tool(input)").unwrap_err();
        assert!(matches!(err, AgentError::ToolNotFound { .. }));
    }

    #[test]
    fn test_loop_runner_history_grows_with_steps() {
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Thought: thinking").unwrap();
        runner.step("Thought: still thinking").unwrap();
        assert!(runner.history().len() >= 2);
    }

    #[test]
    fn test_loop_runner_action_adds_observation_to_history() {
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Action: echo(hi)").unwrap();
        // Two messages: the Action response + the Observation
        assert!(runner.history().len() >= 2);
    }

    #[test]
    fn test_loop_runner_iteration_count_counts_steps_not_messages() {
        // 0.1.x returned the history length, which grows by two per action.
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Thought: one").unwrap();
        runner.step("Action: echo(x)").unwrap();
        assert_eq!(runner.iteration_count(), 2);
        assert_eq!(runner.history().len(), 3);
    }

    // ── 0.2.0 regression tests ───────────────────────────────────────────────

    #[test]
    fn test_parse_thought_then_action_runs_the_action() {
        // 0.1.x matched only the first line's prefix, so the usual model reply
        // "Thought: ...\nAction: ..." was taken as a Thought and the tool never ran.
        let step = parse_react_step("Thought: I should look it up.\nAction: search(rust wasm)").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "rust wasm".into() });
    }

    #[test]
    fn test_parse_final_answer_after_thought() {
        let step = parse_react_step("Thought: I know this now.\nFinal Answer: Paris").unwrap();
        assert_eq!(step, ReActStep::FinalAnswer("Paris".into()));
    }

    #[test]
    fn test_parse_keeps_inner_parentheses() {
        // 0.1.x trimmed every trailing ')' and turned "f(x)" into "f(x".
        let step = parse_react_step("Action: calc(max(1, 2))").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "calc".into(), input: "max(1, 2)".into() });
    }

    #[test]
    fn test_parse_langchain_action_input() {
        let step = parse_react_step("Thought: search\nAction: search\nAction Input: rust crates").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "rust crates".into() });
    }

    #[test]
    fn test_parse_ignores_hallucinated_observation() {
        let text = "Action: search(weather)\nObservation: sunny\nFinal Answer: it is sunny";
        let step = parse_react_step(text).unwrap();
        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "weather".into() });
    }

    #[test]
    fn test_parse_markdown_and_case() {
        let step = parse_react_step("**Action:** echo(hi)").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "echo".into(), input: "hi".into() });
        let step = parse_react_step("final answer: 42").unwrap();
        assert_eq!(step, ReActStep::FinalAnswer("42".into()));
    }

    #[test]
    fn test_parse_colon_input_with_url() {
        let step = parse_react_step("Action: fetch: https://example.com/a").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "fetch".into(), input: "https://example.com/a".into() });
        let step = parse_react_step("Action: fetch(https://example.com/a)").unwrap();
        assert_eq!(step, ReActStep::Action { tool: "fetch".into(), input: "https://example.com/a".into() });
    }

    #[test]
    fn test_parse_empty_action_is_error() {
        assert!(matches!(parse_react_step("Action:   "), Err(AgentError::ParseError(_))));
    }

    #[test]
    fn test_step_enforces_max_iterations() {
        // 0.1.x only enforced max_iterations inside run_scripted.
        let config = AgentConfig { max_iterations: 2, ..Default::default() };
        let reg = ToolRegistry::new();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Thought: 1").unwrap();
        runner.step("Thought: 2").unwrap();
        assert!(matches!(runner.step("Thought: 3"), Err(AgentError::MaxIterationsExceeded(2))));
    }

    #[test]
    fn test_step_at_enforces_timeout() {
        let config = AgentConfig { timeout_ms: Some(1_000), ..Default::default() };
        let reg = ToolRegistry::new();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step_at("Thought: a", 10_000.0).unwrap();
        runner.step_at("Thought: b", 10_900.0).unwrap();
        assert!(matches!(runner.step_at("Thought: c", 11_001.0), Err(AgentError::Timeout { timeout_ms: 1_000 })));
    }

    #[test]
    fn test_failed_tool_is_reported_as_error_observation() {
        let mut reg = ToolRegistry::new();
        reg.register(
            ToolSpec::new("bad", "always fails", "{}"),
            Box::new(|_: &str| ToolResult { tool_name: "bad".into(), output: "no network".into(), success: false }),
        )
        .unwrap();
        let config = AgentConfig::default();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Action: bad(x)").unwrap();
        assert!(runner.transcript().contains("Observation: Error: no network"));
    }

    #[test]
    fn test_transcript_and_system_prompt() {
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg).with_system_prompt();
        runner.ask("say hi");
        runner.step("Thought: use echo\nAction: echo(hi)\nObservation: made up").unwrap();
        let t = runner.transcript();
        assert!(t.contains("- echo: Echoes input"));
        assert!(t.contains("Question: say hi\n"));
        assert!(t.contains("Action: echo(hi)\nObservation: echoed: hi\n"));
        assert!(!t.contains("made up"));
    }

    #[test]
    fn test_recover_from_errors_turns_mistakes_into_observations() {
        let config = AgentConfig { recover_from_errors: true, ..Default::default() };
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        runner.step("Action: serach(x)").unwrap(); // typo in the tool name
        runner.step("I think the answer is 4").unwrap(); // no label
        let t = runner.transcript();
        assert!(t.contains("Observation: Error: there is no tool named 'serach'. Available tools: echo."), "{t}");
        assert!(t.contains("Observation: Error: reply with exactly one of"));
        assert_eq!(runner.run_scripted(&["Action: echo(ok)", "Final Answer: 4"]).unwrap(), "4");
    }

    #[test]
    fn test_step_response_native_round_trip() {
        let config = AgentConfig::default();
        let reg = registry_with_echo();
        let mut runner = LoopRunner::new(&config, &reg);
        let out = runner
            .step_response(r#"{"content":[{"type":"tool_use","id":"toolu_1","name":"echo","input":{"input":"hi"}}]}"#)
            .unwrap();
        match out {
            NativeStep::ToolCalls(r) => {
                assert_eq!(r.len(), 1);
                assert_eq!(r[0].output, "echoed: hi");
                assert_eq!(r[0].id.as_deref(), Some("toolu_1"));
            }
            other => panic!("expected tool calls, got {other:?}"),
        }
        let msgs = runner.history().messages();
        assert_eq!(msgs[0].tool_calls.len(), 1);
        assert_eq!(msgs[1].tool_call_id.as_deref(), Some("toolu_1"));
        let fin = runner.step_response(r#"{"choices":[{"message":{"role":"assistant","content":"done"}}]}"#).unwrap();
        assert_eq!(fin, NativeStep::Final("done".into()));
        assert_eq!(runner.iteration_count(), 2);
    }

    #[test]
    fn test_step_response_unknown_tool() {
        let reg = registry_with_echo();
        let body = r#"{"choices":[{"message":{"tool_calls":[{"id":"c","function":{"name":"nope","arguments":"{}"}}]}}]}"#;
        let strict = AgentConfig::default();
        assert!(matches!(LoopRunner::new(&strict, &reg).step_response(body), Err(AgentError::ToolNotFound { .. })));
        let lenient = AgentConfig { recover_from_errors: true, ..Default::default() };
        match LoopRunner::new(&lenient, &reg).step_response(body).unwrap() {
            NativeStep::ToolCalls(r) => assert!(!r[0].success),
            other => panic!("{other:?}"),
        }
    }
}