Skip to main content

wasm_agent/
types.rs

1// SPDX-License-Identifier: MIT
2//! Core types for the ReAct agent loop.
3
4use serde::{Deserialize, Serialize};
5
6/// A single step in the ReAct loop.
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub enum ReActStep {
9    /// The agent's internal reasoning trace.
10    Thought(String),
11    /// A tool invocation with its name and JSON input.
12    Action {
13        /// Tool name.
14        tool: String,
15        /// Tool input (text, or JSON for object schemas).
16        input: String,
17    },
18    /// The result of a tool invocation.
19    Observation(String),
20    /// The agent's final answer, terminating the loop.
21    FinalAnswer(String),
22}
23
24impl ReActStep {
25    /// Returns a short label for the step kind.
26    pub fn kind(&self) -> &'static str {
27        match self {
28            ReActStep::Thought(_) => "Thought",
29            ReActStep::Action { .. } => "Action",
30            ReActStep::Observation(_) => "Observation",
31            ReActStep::FinalAnswer(_) => "FinalAnswer",
32        }
33    }
34
35    /// Returns `true` if this step terminates the loop.
36    pub fn is_final(&self) -> bool {
37        matches!(self, ReActStep::FinalAnswer(_))
38    }
39
40    /// Returns the primary text content of the step.
41    pub fn content(&self) -> &str {
42        match self {
43            ReActStep::Thought(s) | ReActStep::Observation(s) | ReActStep::FinalAnswer(s) => s,
44            ReActStep::Action { input, .. } => input,
45        }
46    }
47}
48
49/// Role of a participant in the conversation.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51pub enum Role {
52    /// The system prompt that defines the agent's persona and instructions.
53    System,
54    /// A human user turn.
55    User,
56    /// An assistant (model) turn.
57    Assistant,
58    /// A tool result injected into the conversation.
59    Tool,
60}
61
62/// A single message in the conversation history.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct Message {
65    /// Role of the message author.
66    pub role: Role,
67    /// Text content of the message.
68    pub content: String,
69    /// Rough token count estimate (one token per 4 bytes, rounded up).
70    pub token_estimate: usize,
71    /// Native tool calls the model made in this (assistant) message.
72    #[serde(default, skip_serializing_if = "Vec::is_empty")]
73    pub tool_calls: Vec<crate::native::ToolCall>,
74    /// For a tool result: the id of the call it answers.
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub tool_call_id: Option<String>,
77}
78
79impl Message {
80    /// Creates a new message, estimating its token count.
81    pub fn new(role: Role, content: impl Into<String>) -> Self {
82        let content = content.into();
83        // Round up: a short message is never free ("Hi" was 0 tokens in 0.1.x).
84        let token_estimate = content.len().div_ceil(4);
85        Self { role, content, token_estimate, tool_calls: Vec::new(), tool_call_id: None }
86    }
87
88    /// Convenience constructor for system messages.
89    pub fn system(content: impl Into<String>) -> Self { Self::new(Role::System, content) }
90    /// Convenience constructor for user messages.
91    pub fn user(content: impl Into<String>) -> Self { Self::new(Role::User, content) }
92    /// Convenience constructor for assistant messages.
93    pub fn assistant(content: impl Into<String>) -> Self { Self::new(Role::Assistant, content) }
94}
95
96/// Agent configuration.
97#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct AgentConfig {
99    /// Maximum number of ReAct iterations before giving up.
100    pub max_iterations: u32,
101    /// Maximum total tokens allowed in the conversation history.
102    pub context_token_limit: usize,
103    /// Optional wall-clock timeout in milliseconds.
104    pub timeout_ms: Option<u64>,
105    /// Model identifier string.
106    pub model: String,
107    /// When true, an unknown tool or an unparseable reply becomes an
108    /// `Error: ...` observation the model can correct, instead of an error
109    /// that ends the loop. Default false (0.1 behaviour).
110    #[serde(default)]
111    pub recover_from_errors: bool,
112}
113
114impl Default for AgentConfig {
115    fn default() -> Self {
116        Self {
117            max_iterations: 10,
118            context_token_limit: 8192,
119            timeout_ms: Some(30_000),
120            model: "claude-haiku-4-5-20251001".into(),
121            recover_from_errors: false,
122        }
123    }
124}
125
126/// The result of executing a tool.
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct ToolResult {
129    /// Name of the tool that produced this result.
130    pub tool_name: String,
131    /// The tool's textual output.
132    pub output: String,
133    /// Whether execution succeeded.
134    pub success: bool,
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn test_react_step_thought_kind() {
143        let s = ReActStep::Thought("think".into());
144        assert_eq!(s.kind(), "Thought");
145        assert!(!s.is_final());
146    }
147
148    #[test]
149    fn test_react_step_final_answer_is_final() {
150        let s = ReActStep::FinalAnswer("done".into());
151        assert!(s.is_final());
152    }
153
154    #[test]
155    fn test_react_step_action_kind() {
156        let s = ReActStep::Action { tool: "search".into(), input: "query".into() };
157        assert_eq!(s.kind(), "Action");
158    }
159
160    #[test]
161    fn test_react_step_observation_kind() {
162        let s = ReActStep::Observation("result".into());
163        assert_eq!(s.kind(), "Observation");
164        assert!(!s.is_final());
165    }
166
167    #[test]
168    fn test_react_step_content_returns_text() {
169        let s = ReActStep::Thought("my thought".into());
170        assert_eq!(s.content(), "my thought");
171    }
172
173    #[test]
174    fn test_react_step_action_content_is_input() {
175        let s = ReActStep::Action { tool: "t".into(), input: "i".into() };
176        assert_eq!(s.content(), "i");
177    }
178
179    #[test]
180    fn test_message_token_estimate_nonzero_for_nonempty() {
181        let m = Message::user("hello world this is a test");
182        assert!(m.token_estimate > 0);
183    }
184
185    #[test]
186    fn test_message_empty_content_zero_tokens() {
187        let m = Message::user("");
188        assert_eq!(m.token_estimate, 0);
189    }
190
191    #[test]
192    fn test_agent_config_default_has_reasonable_limits() {
193        let c = AgentConfig::default();
194        assert!(c.max_iterations > 0);
195        assert!(c.context_token_limit > 0);
196    }
197
198    #[test]
199    fn test_message_system_has_system_role() {
200        let m = Message::system("sys");
201        assert_eq!(m.role, Role::System);
202    }
203
204    #[test]
205    fn test_message_assistant_has_assistant_role() {
206        let m = Message::assistant("asst");
207        assert_eq!(m.role, Role::Assistant);
208    }
209}