wasm_agent/
loop_runner.rs

1// SPDX-License-Identifier: MIT
2//! ReAct loop runner: Thought -> Action -> Observation -> FinalAnswer.
3
4use crate::error::AgentError;
5use crate::history::ConversationHistory;
6use crate::tools::ToolRegistry;
7use crate::types::{AgentConfig, Message, ReActStep, Role};
8use serde::{Deserialize, Serialize};
9
10/// Find `label` at the start of a line (case-insensitive, leading spaces and
11/// `*`/`#` markdown decoration allowed). Returns the byte offset just after the
12/// label's colon.
13fn find_label(text: &str, label: &str) -> Option<(usize, usize)> {
14    let mut offset = 0;
15    for line in text.split_inclusive('\n') {
16        let trimmed = line.trim_start_matches([' ', '\t', '*', '#']);
17        let lead = line.len() - trimmed.len();
18        if trimmed.len() >= label.len()
19            && trimmed.is_char_boundary(label.len())
20            && trimmed[..label.len()].eq_ignore_ascii_case(label)
21        {
22            let after = trimmed[label.len()..].trim_start_matches('*');
23            if let Some(rest) = after.strip_prefix(':') {
24                let rest = rest.trim_start_matches('*');
25                let start = offset + lead;
26                let content = offset + line.len() - rest.len();
27                return Some((start, content));
28            }
29        }
30        offset += line.len();
31    }
32    None
33}
34
35/// The text after a label up to the next ReAct label (or the end).
36fn section<'a>(text: &'a str, content_start: usize, stop_labels: &[&str]) -> &'a str {
37    let rest = &text[content_start..];
38    let end = stop_labels
39        .iter()
40        .filter_map(|l| find_label(rest, l).map(|(s, _)| s))
41        .min()
42        .unwrap_or(rest.len());
43    rest[..end].trim()
44}
45
46const LABELS: &[&str] = &["Thought", "Action Input", "Action", "Observation", "Final Answer"];
47
48/// Parses a raw LLM response into a [`ReActStep`].
49///
50/// Models usually write several labelled lines in one response, for example
51/// `Thought: ...` followed by `Action: ...`. The most decisive one wins:
52/// `Final Answer` over `Action` over `Thought`. Labels must start a line and
53/// are matched case-insensitively (`**Action:**` also works). Anything from an
54/// `Observation:` line on is ignored, because the model is inventing a tool
55/// result it has not received.
56///
57/// Supported action forms:
58/// - `Action: tool(input)` (one trailing `)` is removed, inner parentheses are kept)
59/// - `Action: tool: input`
60/// - `Action: tool` followed by an `Action Input: input` line (LangChain style)
61/// - `Action: tool` with no input
62///
63/// # Errors
64/// Returns [`AgentError::ParseError`] if no label is found, or `Action:` names no tool.
65pub fn parse_react_step(response: &str) -> Result<ReActStep, AgentError> {
66    // Drop a hallucinated observation and everything after it.
67    let text = match find_label(response, "Observation") {
68        Some((start, _)) => &response[..start],
69        None => response,
70    };
71
72    if let Some((_, c)) = find_label(text, "Final Answer") {
73        return Ok(ReActStep::FinalAnswer(text[c..].trim().to_string()));
74    }
75
76    if let Some((_, c)) = find_label(text, "Action") {
77        // "Action Input:" also starts with "Action"; find_label requires the
78        // colon right after the label, so it never matches here.
79        let rest = section(text, c, LABELS);
80        let explicit_input = find_label(text, "Action Input").map(|(_, ic)| section(text, ic, LABELS));
81        let (tool, input) = split_action(rest);
82        if tool.is_empty() {
83            return Err(AgentError::ParseError(format!("Action without a tool name: {}", response.trim())));
84        }
85        let input = match explicit_input {
86            Some(i) if input.is_empty() => i.to_string(),
87            _ => input,
88        };
89        return Ok(ReActStep::Action { tool, input });
90    }
91
92    if let Some((_, c)) = find_label(text, "Thought") {
93        return Ok(ReActStep::Thought(section(text, c, LABELS).to_string()));
94    }
95
96    Err(AgentError::ParseError(format!("Could not parse ReAct step from: {}", response.trim())))
97}
98
99fn split_action(rest: &str) -> (String, String) {
100    let first_line = rest.lines().next().unwrap_or("").trim();
101    let paren = first_line.find('(');
102    let colon = first_line.find(':');
103    match (paren, colon) {
104        (Some(p), c) if c.is_none_or(|c| p < c) => {
105            let tool = first_line[..p].trim().to_string();
106            // Input may span lines: take everything after "(" in the section.
107            let after = rest[rest.find('(').map_or(0, |i| i + 1)..].trim_end();
108            let input = after.strip_suffix(')').unwrap_or(after).trim().to_string();
109            (tool, input)
110        }
111        (_, Some(c)) => {
112            let tool = first_line[..c].trim().to_string();
113            let after_colon = &rest[rest.find(':').map_or(rest.len(), |i| i + 1)..];
114            (tool, after_colon.trim().to_string())
115        }
116        _ => (first_line.trim_matches('`').to_string(), rest[first_line.len().min(rest.len())..].trim().to_string()),
117    }
118}
119
120/// Loop state shared by [`LoopRunner`] (borrows its config and tools) and
121/// [`Agent`] (owns them).
122#[derive(Debug)]
123struct State {
124    history: ConversationHistory,
125    steps: u32,
126    started_at_ms: Option<f64>,
127}
128
129impl State {
130    fn new(config: &AgentConfig) -> Self {
131        Self { history: ConversationHistory::new(config.context_token_limit), steps: 0, started_at_ms: None }
132    }
133
134    fn check_budget(&mut self, config: &AgentConfig, now_ms: Option<f64>) -> Result<(), AgentError> {
135        if self.steps >= config.max_iterations {
136            return Err(AgentError::MaxIterationsExceeded(config.max_iterations));
137        }
138        if let (Some(now), Some(limit)) = (now_ms, config.timeout_ms) {
139            let start = *self.started_at_ms.get_or_insert(now);
140            if now - start > limit as f64 {
141                return Err(AgentError::Timeout { timeout_ms: limit });
142            }
143        }
144        Ok(())
145    }
146
147    fn observation(result: crate::types::ToolResult) -> String {
148        if result.success { result.output } else { format!("Error: {}", result.output) }
149    }
150
151    fn step(&mut self, config: &AgentConfig, registry: &ToolRegistry, llm_response: &str, now_ms: Option<f64>) -> Result<ReActStep, AgentError> {
152        self.check_budget(config, now_ms)?;
153        let parsed = parse_react_step(llm_response);
154        let (step, obs) = match parsed {
155            Ok(step) => {
156                let obs = match &step {
157                    ReActStep::Action { tool, input } => match registry.dispatch(tool, input) {
158                        Ok(result) => Some(Self::observation(result)),
159                        Err(AgentError::ToolNotFound { name }) if config.recover_from_errors => Some(format!(
160                            "Error: there is no tool named '{name}'. Available tools: {}.",
161                            registry.tool_names().join(", ")
162                        )),
163                        Err(e) => return Err(e),
164                    },
165                    _ => None,
166                };
167                (step, obs)
168            }
169            Err(AgentError::ParseError(_)) if config.recover_from_errors => (
170                ReActStep::Observation("format error".into()),
171                Some("Error: reply with exactly one of 'Thought: ...', 'Action: tool(input)' or 'Final Answer: ...'.".into()),
172            ),
173            Err(e) => return Err(e),
174        };
175        self.steps += 1;
176        // Keep only the part the parser used: a hallucinated "Observation:"
177        // must not end up in the history as if it were real.
178        let kept = match find_label(llm_response, "Observation") {
179            Some((start, _)) => llm_response[..start].trim_end(),
180            None => llm_response,
181        };
182        self.history.push_with_eviction(Message::assistant(kept));
183        if let Some(obs) = obs {
184            self.history.push_with_eviction(Message::new(Role::Tool, obs));
185        }
186        Ok(step)
187    }
188
189    fn step_response(
190        &mut self,
191        config: &AgentConfig,
192        registry: &ToolRegistry,
193        response_json: &str,
194        now_ms: Option<f64>,
195    ) -> Result<NativeStep, AgentError> {
196        self.check_budget(config, now_ms)?;
197        let turn = crate::native::parse_model_response(response_json)?;
198        self.steps += 1;
199        let text = turn.text.clone().unwrap_or_default();
200        if turn.calls.is_empty() {
201            self.history.push_with_eviction(Message::assistant(text.clone()));
202            return Ok(NativeStep::Final(text));
203        }
204        let mut assistant = Message::assistant(text);
205        assistant.tool_calls = turn.calls.clone();
206        self.history.push_with_eviction(assistant);
207        let mut outputs = Vec::with_capacity(turn.calls.len());
208        for call in &turn.calls {
209            let output = match registry.dispatch_call(call) {
210                Ok(result) => crate::native::ToolOutput { id: call.id.clone(), name: call.name.clone(), success: result.success, output: result.output },
211                Err(AgentError::ToolNotFound { name }) if config.recover_from_errors => crate::native::ToolOutput {
212                    id: call.id.clone(),
213                    name: call.name.clone(),
214                    success: false,
215                    output: format!("there is no tool named '{name}'"),
216                },
217                Err(e) => return Err(e),
218            };
219            let mut m = Message::new(Role::Tool, Self::observation(crate::types::ToolResult {
220                tool_name: output.name.clone(),
221                output: output.output.clone(),
222                success: output.success,
223            }));
224            m.tool_call_id = output.id.clone();
225            self.history.push_with_eviction(m);
226            outputs.push(output);
227        }
228        Ok(NativeStep::ToolCalls(outputs))
229    }
230
231    fn transcript(&self) -> String {
232        let mut out = String::new();
233        for m in self.history.messages() {
234            match m.role {
235                Role::Tool => {
236                    out.push_str("Observation: ");
237                    out.push_str(&m.content);
238                }
239                Role::User => {
240                    out.push_str("Question: ");
241                    out.push_str(&m.content);
242                }
243                _ => out.push_str(&m.content),
244            }
245            out.push('\n');
246        }
247        out
248    }
249}
250
251/// What one native tool-calling turn did (see [`LoopRunner::step_response`]).
252#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
253pub enum NativeStep {
254    /// The model called tools; here are their results, in call order. Send
255    /// them back to the model (see [`crate::native::tool_results_openai`] and
256    /// [`crate::native::tool_results_anthropic`]) and step again.
257    ToolCalls(Vec<crate::native::ToolOutput>),
258    /// The model answered without calling a tool.
259    Final(String),
260}
261
262/// The ReAct loop execution context.
263///
264/// Holds references to the agent configuration and tool registry and drives
265/// the Thought -> Action -> Observation -> FinalAnswer cycle. It never calls a
266/// model itself: you pass each model response to [`step`](Self::step) (text
267/// ReAct) or [`step_response`](Self::step_response) (native tool calling).
268/// Use [`Agent`] if you would rather not manage the borrows.
269pub struct LoopRunner<'a> {
270    config: &'a AgentConfig,
271    registry: &'a ToolRegistry,
272    state: State,
273}
274
275impl<'a> LoopRunner<'a> {
276    /// Creates a new runner with its own conversation history.
277    pub fn new(config: &'a AgentConfig, registry: &'a ToolRegistry) -> Self {
278        Self { config, registry, state: State::new(config) }
279    }
280
281    /// Adds the standard ReAct system prompt (format rules plus the tool list)
282    /// as the first message. System messages are never evicted from history.
283    pub fn with_system_prompt(mut self) -> Self {
284        self.state.history.push_with_eviction(Message::system(system_prompt(self.registry)));
285        self
286    }
287
288    /// Adds the user's question to the history.
289    pub fn ask(&mut self, question: &str) {
290        self.state.history.push_with_eviction(Message::user(question));
291    }
292
293    /// Processes one text LLM response, dispatching the tool when it is an Action.
294    ///
295    /// With `config.recover_from_errors`, an unknown tool or an unparseable
296    /// reply becomes an `Error: ...` observation the model can correct,
297    /// instead of an error that ends the loop.
298    ///
299    /// # Errors
300    /// - [`AgentError::MaxIterationsExceeded`]: more than `config.max_iterations` steps.
301    /// - [`AgentError::ParseError`]: the response does not match any ReAct format.
302    /// - [`AgentError::ToolNotFound`]: the requested tool is not in the registry.
303    pub fn step(&mut self, llm_response: &str) -> Result<ReActStep, AgentError> {
304        self.state.step(self.config, self.registry, llm_response, None)
305    }
306
307    /// Like [`step`](Self::step), and also enforces `config.timeout_ms`. Pass
308    /// the current time in milliseconds (for example `Date.now()`); the first
309    /// call starts the clock.
310    ///
311    /// # Errors
312    /// [`AgentError::Timeout`] once more than `timeout_ms` has passed since the
313    /// first call, plus everything [`step`](Self::step) returns.
314    pub fn step_at(&mut self, llm_response: &str, now_ms: f64) -> Result<ReActStep, AgentError> {
315        self.state.step(self.config, self.registry, llm_response, Some(now_ms))
316    }
317
318    /// Processes one native tool-calling response: the raw JSON body from
319    /// OpenAI Chat Completions or Responses, Anthropic Messages, or Ollama's
320    /// chat API. Every tool call in it is dispatched (arguments arrive as JSON;
321    /// see [`ToolRegistry::dispatch_call`]) and the results are returned and
322    /// kept in the history with their call ids.
323    ///
324    /// # Errors
325    /// [`AgentError::ParseError`] if the body is not a recognised response,
326    /// [`AgentError::MaxIterationsExceeded`], [`AgentError::ToolNotFound`]
327    /// (unless `recover_from_errors`).
328    pub fn step_response(&mut self, response_json: &str) -> Result<NativeStep, AgentError> {
329        self.state.step_response(self.config, self.registry, response_json, None)
330    }
331
332    /// Runs a pre-scripted sequence of LLM responses, returning the final answer.
333    ///
334    /// Useful for offline simulation and deterministic testing.
335    ///
336    /// # Errors
337    /// - [`AgentError::MaxIterationsExceeded`]: sequence exhausted before a `FinalAnswer`.
338    /// - Any error from [`step`](Self::step).
339    pub fn run_scripted(&mut self, responses: &[&str]) -> Result<String, AgentError> {
340        for response in responses {
341            if let ReActStep::FinalAnswer(answer) = self.step(response)? {
342                return Ok(answer);
343            }
344        }
345        Err(AgentError::MaxIterationsExceeded(self.config.max_iterations))
346    }
347
348    /// The conversation so far as one prompt string, ready to send to a
349    /// completion-style model: system and user messages as-is, model turns
350    /// as-is, tool results as `Observation: ...` lines.
351    pub fn transcript(&self) -> String {
352        self.state.transcript()
353    }
354
355    /// Returns a reference to the accumulated conversation history.
356    pub fn history(&self) -> &ConversationHistory { &self.state.history }
357
358    /// Number of model responses processed so far.
359    pub fn iteration_count(&self) -> usize { self.state.steps as usize }
360}
361
362/// A ReAct agent that owns its configuration, tools and history (no
363/// lifetimes to manage). Same behaviour as [`LoopRunner`].
364///
365/// ```
366/// use wasm_agent::{Agent, AgentConfig, ReActStep, ToolSpec};
367/// use wasm_agent::types::ToolResult;
368///
369/// let mut agent = Agent::new(AgentConfig::default());
370/// agent.tools_mut().register(
371///     ToolSpec::new("echo", "Repeats its input", r#"{"type":"string"}"#),
372///     Box::new(|s: &str| ToolResult { tool_name: "echo".into(), output: s.into(), success: true }),
373/// )?;
374/// agent.step("Action: echo(hi)")?;
375/// assert_eq!(agent.step("Final Answer: done")?, ReActStep::FinalAnswer("done".into()));
376/// # Ok::<(), wasm_agent::AgentError>(())
377/// ```
378pub struct Agent {
379    config: AgentConfig,
380    registry: ToolRegistry,
381    state: State,
382}
383
384impl Agent {
385    /// An agent with no tools yet.
386    pub fn new(config: AgentConfig) -> Self {
387        let state = State::new(&config);
388        Self { config, registry: ToolRegistry::new(), state }
389    }
390
391    /// An agent with the given tools.
392    pub fn with_tools(config: AgentConfig, registry: ToolRegistry) -> Self {
393        let state = State::new(&config);
394        Self { config, registry, state }
395    }
396
397    /// The tool registry, to register more tools.
398    pub fn tools_mut(&mut self) -> &mut ToolRegistry { &mut self.registry }
399
400    /// The tool registry.
401    pub fn tools(&self) -> &ToolRegistry { &self.registry }
402
403    /// The configuration.
404    pub fn config(&self) -> &AgentConfig { &self.config }
405
406    /// Adds the standard ReAct system prompt as the first message.
407    pub fn add_system_prompt(&mut self) {
408        let prompt = system_prompt(&self.registry);
409        self.state.history.push_with_eviction(Message::system(prompt));
410    }
411
412    /// Adds the user's question to the history.
413    pub fn ask(&mut self, question: &str) {
414        self.state.history.push_with_eviction(Message::user(question));
415    }
416
417    /// See [`LoopRunner::step`].
418    ///
419    /// # Errors
420    /// See [`LoopRunner::step`].
421    pub fn step(&mut self, llm_response: &str) -> Result<ReActStep, AgentError> {
422        self.state.step(&self.config, &self.registry, llm_response, None)
423    }
424
425    /// See [`LoopRunner::step_at`].
426    ///
427    /// # Errors
428    /// See [`LoopRunner::step_at`].
429    pub fn step_at(&mut self, llm_response: &str, now_ms: f64) -> Result<ReActStep, AgentError> {
430        self.state.step(&self.config, &self.registry, llm_response, Some(now_ms))
431    }
432
433    /// See [`LoopRunner::step_response`].
434    ///
435    /// # Errors
436    /// See [`LoopRunner::step_response`].
437    pub fn step_response(&mut self, response_json: &str) -> Result<NativeStep, AgentError> {
438        self.state.step_response(&self.config, &self.registry, response_json, None)
439    }
440
441    /// See [`LoopRunner::transcript`].
442    pub fn transcript(&self) -> String { self.state.transcript() }
443
444    /// The conversation history.
445    pub fn history(&self) -> &ConversationHistory { &self.state.history }
446
447    /// Number of model responses processed so far.
448    pub fn iteration_count(&self) -> usize { self.state.steps as usize }
449}
450
451/// The standard ReAct instructions plus the registry's tool list, for use as
452/// a system prompt.
453pub fn system_prompt(registry: &ToolRegistry) -> String {
454    format!(
455        "Answer the question. You can use tools.\n\n{}\n\n\
456         Use this format, one step per reply:\n\
457         Thought: what you are thinking\n\
458         Action: tool_name(input)\n\
459         Then wait: the tool result comes back as \"Observation: ...\".\n\
460         When you know the answer, reply with:\n\
461         Final Answer: the answer",
462        registry.tools_prompt()
463    )
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469    use crate::tools::{ToolRegistry, ToolSpec};
470    use crate::types::ToolResult;
471
472    fn registry_with_echo() -> ToolRegistry {
473        let mut reg = ToolRegistry::new();
474        reg.register(
475            ToolSpec::new("echo", "Echoes input", "{}"),
476            Box::new(|input: &str| ToolResult {
477                tool_name: "echo".into(),
478                output: format!("echoed: {input}"),
479                success: true,
480            }),
481        ).unwrap();
482        reg
483    }
484
485    #[test]
486    fn test_parse_thought() {
487        let step = parse_react_step("Thought: I need to search for information").unwrap();
488        assert!(matches!(step, ReActStep::Thought(_)));
489    }
490
491    #[test]
492    fn test_parse_action_paren_syntax() {
493        let step = parse_react_step("Action: echo(hello world)").unwrap();
494        assert!(matches!(&step, ReActStep::Action { tool, input }
495            if tool == "echo" && input == "hello world"));
496    }
497
498    #[test]
499    fn test_parse_action_colon_syntax() {
500        let step = parse_react_step("Action: echo: hello").unwrap();
501        assert!(matches!(&step, ReActStep::Action { tool, .. } if tool == "echo"));
502    }
503
504    #[test]
505    fn test_parse_final_answer() {
506        let step = parse_react_step("Final Answer: 42").unwrap();
507        assert!(matches!(&step, ReActStep::FinalAnswer(s) if s == "42"));
508    }
509
510    #[test]
511    fn test_parse_unknown_format_returns_parse_error() {
512        let err = parse_react_step("This is just random text without a prefix").unwrap_err();
513        assert!(matches!(err, AgentError::ParseError(_)));
514    }
515
516    #[test]
517    fn test_parse_action_no_args() {
518        let step = parse_react_step("Action: mytool").unwrap();
519        assert!(matches!(&step, ReActStep::Action { tool, .. } if tool == "mytool"));
520    }
521
522    #[test]
523    fn test_loop_runner_scripted_final_answer() {
524        let config = AgentConfig::default();
525        let reg = registry_with_echo();
526        let mut runner = LoopRunner::new(&config, &reg);
527        let answer = runner.run_scripted(&[
528            "Thought: Let me think",
529            "Action: echo(test)",
530            "Final Answer: done",
531        ]).unwrap();
532        assert_eq!(answer, "done");
533    }
534
535    #[test]
536    fn test_loop_runner_max_iterations_returns_error() {
537        let config = AgentConfig { max_iterations: 2, ..Default::default() };
538        let reg = ToolRegistry::new();
539        let mut runner = LoopRunner::new(&config, &reg);
540        let err = runner.run_scripted(&[
541            "Thought: step 1",
542            "Thought: step 2",
543            "Thought: step 3",
544        ]).unwrap_err();
545        assert!(matches!(err, AgentError::MaxIterationsExceeded(2)));
546    }
547
548    #[test]
549    fn test_loop_runner_tool_not_found_returns_error() {
550        let config = AgentConfig::default();
551        let reg = ToolRegistry::new();
552        let mut runner = LoopRunner::new(&config, &reg);
553        let err = runner.step("Action: missing_tool(input)").unwrap_err();
554        assert!(matches!(err, AgentError::ToolNotFound { .. }));
555    }
556
557    #[test]
558    fn test_loop_runner_history_grows_with_steps() {
559        let config = AgentConfig::default();
560        let reg = registry_with_echo();
561        let mut runner = LoopRunner::new(&config, &reg);
562        runner.step("Thought: thinking").unwrap();
563        runner.step("Thought: still thinking").unwrap();
564        assert!(runner.history().len() >= 2);
565    }
566
567    #[test]
568    fn test_loop_runner_action_adds_observation_to_history() {
569        let config = AgentConfig::default();
570        let reg = registry_with_echo();
571        let mut runner = LoopRunner::new(&config, &reg);
572        runner.step("Action: echo(hi)").unwrap();
573        // Two messages: the Action response + the Observation
574        assert!(runner.history().len() >= 2);
575    }
576
577    #[test]
578    fn test_loop_runner_iteration_count_counts_steps_not_messages() {
579        // 0.1.x returned the history length, which grows by two per action.
580        let config = AgentConfig::default();
581        let reg = registry_with_echo();
582        let mut runner = LoopRunner::new(&config, &reg);
583        runner.step("Thought: one").unwrap();
584        runner.step("Action: echo(x)").unwrap();
585        assert_eq!(runner.iteration_count(), 2);
586        assert_eq!(runner.history().len(), 3);
587    }
588
589    // ── 0.2.0 regression tests ───────────────────────────────────────────────
590
591    #[test]
592    fn test_parse_thought_then_action_runs_the_action() {
593        // 0.1.x matched only the first line's prefix, so the usual model reply
594        // "Thought: ...\nAction: ..." was taken as a Thought and the tool never ran.
595        let step = parse_react_step("Thought: I should look it up.\nAction: search(rust wasm)").unwrap();
596        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "rust wasm".into() });
597    }
598
599    #[test]
600    fn test_parse_final_answer_after_thought() {
601        let step = parse_react_step("Thought: I know this now.\nFinal Answer: Paris").unwrap();
602        assert_eq!(step, ReActStep::FinalAnswer("Paris".into()));
603    }
604
605    #[test]
606    fn test_parse_keeps_inner_parentheses() {
607        // 0.1.x trimmed every trailing ')' and turned "f(x)" into "f(x".
608        let step = parse_react_step("Action: calc(max(1, 2))").unwrap();
609        assert_eq!(step, ReActStep::Action { tool: "calc".into(), input: "max(1, 2)".into() });
610    }
611
612    #[test]
613    fn test_parse_langchain_action_input() {
614        let step = parse_react_step("Thought: search\nAction: search\nAction Input: rust crates").unwrap();
615        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "rust crates".into() });
616    }
617
618    #[test]
619    fn test_parse_ignores_hallucinated_observation() {
620        let text = "Action: search(weather)\nObservation: sunny\nFinal Answer: it is sunny";
621        let step = parse_react_step(text).unwrap();
622        assert_eq!(step, ReActStep::Action { tool: "search".into(), input: "weather".into() });
623    }
624
625    #[test]
626    fn test_parse_markdown_and_case() {
627        let step = parse_react_step("**Action:** echo(hi)").unwrap();
628        assert_eq!(step, ReActStep::Action { tool: "echo".into(), input: "hi".into() });
629        let step = parse_react_step("final answer: 42").unwrap();
630        assert_eq!(step, ReActStep::FinalAnswer("42".into()));
631    }
632
633    #[test]
634    fn test_parse_colon_input_with_url() {
635        let step = parse_react_step("Action: fetch: https://example.com/a").unwrap();
636        assert_eq!(step, ReActStep::Action { tool: "fetch".into(), input: "https://example.com/a".into() });
637        let step = parse_react_step("Action: fetch(https://example.com/a)").unwrap();
638        assert_eq!(step, ReActStep::Action { tool: "fetch".into(), input: "https://example.com/a".into() });
639    }
640
641    #[test]
642    fn test_parse_empty_action_is_error() {
643        assert!(matches!(parse_react_step("Action:   "), Err(AgentError::ParseError(_))));
644    }
645
646    #[test]
647    fn test_step_enforces_max_iterations() {
648        // 0.1.x only enforced max_iterations inside run_scripted.
649        let config = AgentConfig { max_iterations: 2, ..Default::default() };
650        let reg = ToolRegistry::new();
651        let mut runner = LoopRunner::new(&config, &reg);
652        runner.step("Thought: 1").unwrap();
653        runner.step("Thought: 2").unwrap();
654        assert!(matches!(runner.step("Thought: 3"), Err(AgentError::MaxIterationsExceeded(2))));
655    }
656
657    #[test]
658    fn test_step_at_enforces_timeout() {
659        let config = AgentConfig { timeout_ms: Some(1_000), ..Default::default() };
660        let reg = ToolRegistry::new();
661        let mut runner = LoopRunner::new(&config, &reg);
662        runner.step_at("Thought: a", 10_000.0).unwrap();
663        runner.step_at("Thought: b", 10_900.0).unwrap();
664        assert!(matches!(runner.step_at("Thought: c", 11_001.0), Err(AgentError::Timeout { timeout_ms: 1_000 })));
665    }
666
667    #[test]
668    fn test_failed_tool_is_reported_as_error_observation() {
669        let mut reg = ToolRegistry::new();
670        reg.register(
671            ToolSpec::new("bad", "always fails", "{}"),
672            Box::new(|_: &str| ToolResult { tool_name: "bad".into(), output: "no network".into(), success: false }),
673        )
674        .unwrap();
675        let config = AgentConfig::default();
676        let mut runner = LoopRunner::new(&config, &reg);
677        runner.step("Action: bad(x)").unwrap();
678        assert!(runner.transcript().contains("Observation: Error: no network"));
679    }
680
681    #[test]
682    fn test_transcript_and_system_prompt() {
683        let config = AgentConfig::default();
684        let reg = registry_with_echo();
685        let mut runner = LoopRunner::new(&config, &reg).with_system_prompt();
686        runner.ask("say hi");
687        runner.step("Thought: use echo\nAction: echo(hi)\nObservation: made up").unwrap();
688        let t = runner.transcript();
689        assert!(t.contains("- echo: Echoes input"));
690        assert!(t.contains("Question: say hi\n"));
691        assert!(t.contains("Action: echo(hi)\nObservation: echoed: hi\n"));
692        assert!(!t.contains("made up"));
693    }
694
695    #[test]
696    fn test_recover_from_errors_turns_mistakes_into_observations() {
697        let config = AgentConfig { recover_from_errors: true, ..Default::default() };
698        let reg = registry_with_echo();
699        let mut runner = LoopRunner::new(&config, &reg);
700        runner.step("Action: serach(x)").unwrap(); // typo in the tool name
701        runner.step("I think the answer is 4").unwrap(); // no label
702        let t = runner.transcript();
703        assert!(t.contains("Observation: Error: there is no tool named 'serach'. Available tools: echo."), "{t}");
704        assert!(t.contains("Observation: Error: reply with exactly one of"));
705        assert_eq!(runner.run_scripted(&["Action: echo(ok)", "Final Answer: 4"]).unwrap(), "4");
706    }
707
708    #[test]
709    fn test_step_response_native_round_trip() {
710        let config = AgentConfig::default();
711        let reg = registry_with_echo();
712        let mut runner = LoopRunner::new(&config, &reg);
713        let out = runner
714            .step_response(r#"{"content":[{"type":"tool_use","id":"toolu_1","name":"echo","input":{"input":"hi"}}]}"#)
715            .unwrap();
716        match out {
717            NativeStep::ToolCalls(r) => {
718                assert_eq!(r.len(), 1);
719                assert_eq!(r[0].output, "echoed: hi");
720                assert_eq!(r[0].id.as_deref(), Some("toolu_1"));
721            }
722            other => panic!("expected tool calls, got {other:?}"),
723        }
724        let msgs = runner.history().messages();
725        assert_eq!(msgs[0].tool_calls.len(), 1);
726        assert_eq!(msgs[1].tool_call_id.as_deref(), Some("toolu_1"));
727        let fin = runner.step_response(r#"{"choices":[{"message":{"role":"assistant","content":"done"}}]}"#).unwrap();
728        assert_eq!(fin, NativeStep::Final("done".into()));
729        assert_eq!(runner.iteration_count(), 2);
730    }
731
732    #[test]
733    fn test_step_response_unknown_tool() {
734        let reg = registry_with_echo();
735        let body = r#"{"choices":[{"message":{"tool_calls":[{"id":"c","function":{"name":"nope","arguments":"{}"}}]}}]}"#;
736        let strict = AgentConfig::default();
737        assert!(matches!(LoopRunner::new(&strict, &reg).step_response(body), Err(AgentError::ToolNotFound { .. })));
738        let lenient = AgentConfig { recover_from_errors: true, ..Default::default() };
739        match LoopRunner::new(&lenient, &reg).step_response(body).unwrap() {
740            NativeStep::ToolCalls(r) => assert!(!r[0].success),
741            other => panic!("{other:?}"),
742        }
743    }
744}
wasm_agent/loop_runner.rs

wasm_agent/
loop_runner.rs
