use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PlanningModeConfig {
pub enabled: bool,
pub max_planned_steps: usize,
pub replan_on_failure: bool,
pub max_replan_attempts: usize,
pub include_reasoning: bool,
pub auto_execute_threshold: f64,
}
impl Default for PlanningModeConfig {
fn default() -> Self {
Self {
enabled: false,
max_planned_steps: 10,
replan_on_failure: true,
max_replan_attempts: 3,
include_reasoning: true,
auto_execute_threshold: 0.8,
}
}
}
impl PlanningModeConfig {
pub fn enabled() -> Self {
Self {
enabled: true,
..Default::default()
}
}
pub fn with_max_steps(mut self, max: usize) -> Self {
self.max_planned_steps = max;
self
}
pub fn with_replan(mut self, enabled: bool) -> Self {
self.replan_on_failure = enabled;
self
}
pub fn with_auto_execute_threshold(mut self, threshold: f64) -> Self {
self.auto_execute_threshold = threshold.clamp(0.0, 1.0);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlannedStep {
pub id: String,
pub action: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub checkpoint: Option<Checkpoint>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub depends_on: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(default = "default_confidence")]
pub confidence: f64,
#[serde(default = "default_true")]
pub critical: bool,
}
fn default_confidence() -> f64 {
0.5
}
fn default_true() -> bool {
true
}
impl PlannedStep {
pub fn new(id: impl Into<String>, action: Value) -> Self {
Self {
id: id.into(),
action,
checkpoint: None,
depends_on: Vec::new(),
description: None,
confidence: 0.5,
critical: true,
}
}
pub fn with_checkpoint(mut self, checkpoint: Checkpoint) -> Self {
self.checkpoint = Some(checkpoint);
self
}
pub fn depends_on(mut self, step_id: impl Into<String>) -> Self {
self.depends_on.push(step_id.into());
self
}
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = Some(desc.into());
self
}
pub fn with_confidence(mut self, confidence: f64) -> Self {
self.confidence = confidence.clamp(0.0, 1.0);
self
}
pub fn with_critical(mut self, critical: bool) -> Self {
self.critical = critical;
self
}
pub fn from_json(value: &Value) -> Option<Self> {
let id = value
.get("id")
.and_then(|v| v.as_str())
.unwrap_or("step")
.to_string();
let action = value.get("action")?.clone();
let checkpoint = value.get("checkpoint").and_then(Checkpoint::from_json);
let depends_on = value
.get("depends_on")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let description = value
.get("description")
.and_then(|v| v.as_str())
.map(String::from);
let confidence = value
.get("confidence")
.and_then(|v| v.as_f64())
.unwrap_or(0.5);
let critical = value
.get("critical")
.and_then(|v| v.as_bool())
.unwrap_or(true);
Some(Self {
id,
action,
checkpoint,
depends_on,
description,
confidence: confidence.clamp(0.0, 1.0),
critical,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Checkpoint {
pub checkpoint_type: CheckpointType,
pub expected: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
}
impl Checkpoint {
pub fn url_contains(pattern: impl Into<String>) -> Self {
Self {
checkpoint_type: CheckpointType::UrlContains,
expected: pattern.into(),
timeout_ms: None,
description: None,
}
}
pub fn element_exists(selector: impl Into<String>) -> Self {
Self {
checkpoint_type: CheckpointType::ElementExists,
expected: selector.into(),
timeout_ms: Some(5000),
description: None,
}
}
pub fn text_contains(text: impl Into<String>) -> Self {
Self {
checkpoint_type: CheckpointType::TextContains,
expected: text.into(),
timeout_ms: None,
description: None,
}
}
pub fn js_condition(condition: impl Into<String>) -> Self {
Self {
checkpoint_type: CheckpointType::JsCondition,
expected: condition.into(),
timeout_ms: None,
description: None,
}
}
pub fn element_not_exists(selector: impl Into<String>) -> Self {
Self {
checkpoint_type: CheckpointType::ElementNotExists,
expected: selector.into(),
timeout_ms: Some(5000),
description: None,
}
}
pub fn with_timeout(mut self, ms: u64) -> Self {
self.timeout_ms = Some(ms);
self
}
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = Some(desc.into());
self
}
pub fn from_json(value: &Value) -> Option<Self> {
let type_str = value.get("type").and_then(|v| v.as_str())?;
let expected = value.get("expected").and_then(|v| v.as_str())?.to_string();
let checkpoint_type = match type_str {
"url_contains" => CheckpointType::UrlContains,
"element_exists" => CheckpointType::ElementExists,
"element_not_exists" => CheckpointType::ElementNotExists,
"text_contains" => CheckpointType::TextContains,
"js_condition" => CheckpointType::JsCondition,
"page_loaded" => CheckpointType::PageLoaded,
_ => return None,
};
let timeout_ms = value.get("timeout_ms").and_then(|v| v.as_u64());
let description = value
.get("description")
.and_then(|v| v.as_str())
.map(String::from);
Some(Self {
checkpoint_type,
expected,
timeout_ms,
description,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CheckpointType {
UrlContains,
ElementExists,
ElementNotExists,
TextContains,
JsCondition,
PageLoaded,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionPlan {
pub goal: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
pub steps: Vec<PlannedStep>,
pub confidence: f64,
pub is_complete: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub expected_outcome: Option<String>,
}
impl ExecutionPlan {
pub fn new(goal: impl Into<String>) -> Self {
Self {
goal: goal.into(),
reasoning: None,
steps: Vec::new(),
confidence: 0.5,
is_complete: false,
expected_outcome: None,
}
}
pub fn add_step(mut self, step: PlannedStep) -> Self {
self.steps.push(step);
self
}
pub fn with_reasoning(mut self, reasoning: impl Into<String>) -> Self {
self.reasoning = Some(reasoning.into());
self
}
pub fn with_confidence(mut self, confidence: f64) -> Self {
self.confidence = confidence.clamp(0.0, 1.0);
self
}
pub fn complete(mut self) -> Self {
self.is_complete = true;
self
}
pub fn with_expected_outcome(mut self, outcome: impl Into<String>) -> Self {
self.expected_outcome = Some(outcome.into());
self
}
pub fn from_json(value: &Value) -> Option<Self> {
let goal = value
.get("goal")
.and_then(|v| v.as_str())
.unwrap_or("automation")
.to_string();
let reasoning = value
.get("reasoning")
.and_then(|v| v.as_str())
.map(String::from);
let steps = value
.get("steps")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(PlannedStep::from_json).collect())
.unwrap_or_default();
let confidence = value
.get("confidence")
.and_then(|v| v.as_f64())
.unwrap_or(0.5);
let is_complete = value
.get("is_complete")
.or_else(|| value.get("done"))
.and_then(|v| v.as_bool())
.unwrap_or(false);
let expected_outcome = value
.get("expected_outcome")
.and_then(|v| v.as_str())
.map(String::from);
Some(Self {
goal,
reasoning,
steps,
confidence: confidence.clamp(0.0, 1.0),
is_complete,
expected_outcome,
})
}
pub fn step_count(&self) -> usize {
self.steps.len()
}
pub fn is_empty(&self) -> bool {
self.steps.is_empty()
}
pub fn average_step_confidence(&self) -> f64 {
if self.steps.is_empty() {
0.0
} else {
let sum: f64 = self.steps.iter().map(|s| s.confidence).sum();
sum / self.steps.len() as f64
}
}
pub fn steps_with_checkpoints(&self) -> Vec<&PlannedStep> {
self.steps
.iter()
.filter(|s| s.checkpoint.is_some())
.collect()
}
pub fn critical_steps(&self) -> Vec<&PlannedStep> {
self.steps.iter().filter(|s| s.critical).collect()
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PlanExecutionState {
pub completed_steps: Vec<String>,
pub failed_steps: Vec<(String, String)>,
pub current_step: usize,
pub replan_attempts: usize,
pub should_continue: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_checkpoint_result: Option<CheckpointResult>,
}
impl PlanExecutionState {
pub fn new() -> Self {
Self {
completed_steps: Vec::new(),
failed_steps: Vec::new(),
current_step: 0,
replan_attempts: 0,
should_continue: true,
last_checkpoint_result: None,
}
}
pub fn complete_step(&mut self, step_id: impl Into<String>) {
self.completed_steps.push(step_id.into());
self.current_step += 1;
}
pub fn fail_step(&mut self, step_id: impl Into<String>, error: impl Into<String>) {
self.failed_steps.push((step_id.into(), error.into()));
}
pub fn is_completed(&self, step_id: &str) -> bool {
self.completed_steps.iter().any(|s| s == step_id)
}
pub fn is_failed(&self, step_id: &str) -> bool {
self.failed_steps.iter().any(|(s, _)| s == step_id)
}
pub fn steps_executed(&self) -> usize {
self.completed_steps.len() + self.failed_steps.len()
}
pub fn has_failures(&self) -> bool {
!self.failed_steps.is_empty()
}
pub fn first_error(&self) -> Option<&str> {
self.failed_steps.first().map(|(_, e)| e.as_str())
}
pub fn stop(&mut self) {
self.should_continue = false;
}
pub fn increment_replan(&mut self) {
self.replan_attempts += 1;
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointResult {
pub passed: bool,
pub checkpoint: Checkpoint,
#[serde(skip_serializing_if = "Option::is_none")]
pub actual_value: Option<String>,
pub duration_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
impl CheckpointResult {
pub fn passed(checkpoint: Checkpoint, duration_ms: u64) -> Self {
Self {
passed: true,
checkpoint,
actual_value: None,
duration_ms,
error: None,
}
}
pub fn failed(checkpoint: Checkpoint, error: impl Into<String>, duration_ms: u64) -> Self {
Self {
passed: false,
checkpoint,
actual_value: None,
duration_ms,
error: Some(error.into()),
}
}
pub fn with_actual(mut self, value: impl Into<String>) -> Self {
self.actual_value = Some(value.into());
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReplanContext {
pub goal: String,
pub completed_steps: Vec<String>,
pub failed_step: String,
pub error: String,
pub current_state: PageState,
#[serde(skip_serializing_if = "Option::is_none")]
pub previous_plan: Option<ExecutionPlan>,
}
impl ReplanContext {
pub fn new(
goal: impl Into<String>,
failed_step: impl Into<String>,
error: impl Into<String>,
current_state: PageState,
) -> Self {
Self {
goal: goal.into(),
completed_steps: Vec::new(),
failed_step: failed_step.into(),
error: error.into(),
current_state,
previous_plan: None,
}
}
pub fn with_completed(mut self, steps: Vec<String>) -> Self {
self.completed_steps = steps;
self
}
pub fn with_previous_plan(mut self, plan: ExecutionPlan) -> Self {
self.previous_plan = Some(plan);
self
}
pub fn to_prompt(&self) -> String {
let mut prompt = String::with_capacity(1024);
prompt.push_str("RE-PLANNING REQUIRED\n\n");
prompt.push_str("Goal: ");
prompt.push_str(&self.goal);
prompt.push_str("\n\n");
if !self.completed_steps.is_empty() {
prompt.push_str("Completed steps:\n");
for step in &self.completed_steps {
prompt.push_str("- ");
prompt.push_str(step);
prompt.push('\n');
}
prompt.push('\n');
}
prompt.push_str("Failed step: ");
prompt.push_str(&self.failed_step);
prompt.push_str("\nError: ");
prompt.push_str(&self.error);
prompt.push_str("\n\n");
prompt.push_str("Current state:\n");
prompt.push_str("- URL: ");
prompt.push_str(&self.current_state.url);
prompt.push_str("\n- Title: ");
prompt.push_str(&self.current_state.title);
prompt.push('\n');
prompt.push_str("\nPlease create a new plan to achieve the goal from the current state.");
prompt
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PageState {
pub url: String,
pub title: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub html_snippet: Option<String>,
}
impl PageState {
pub fn new(url: impl Into<String>, title: impl Into<String>) -> Self {
Self {
url: url.into(),
title: title.into(),
html_snippet: None,
}
}
pub fn with_html(mut self, html: impl Into<String>) -> Self {
self.html_snippet = Some(html.into());
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_planning_mode_config() {
let config = PlanningModeConfig::enabled()
.with_max_steps(5)
.with_replan(false);
assert!(config.enabled);
assert_eq!(config.max_planned_steps, 5);
assert!(!config.replan_on_failure);
}
#[test]
fn test_planned_step() {
let step = PlannedStep::new("step1", serde_json::json!({"Click": "button"}))
.with_checkpoint(Checkpoint::element_exists(".success"))
.with_confidence(0.9)
.depends_on("step0");
assert_eq!(step.id, "step1");
assert!(step.checkpoint.is_some());
assert_eq!(step.confidence, 0.9);
assert_eq!(step.depends_on, vec!["step0"]);
}
#[test]
fn test_checkpoint_types() {
let url_check = Checkpoint::url_contains("/dashboard");
assert_eq!(url_check.checkpoint_type, CheckpointType::UrlContains);
let elem_check = Checkpoint::element_exists(".modal").with_timeout(10000);
assert_eq!(elem_check.checkpoint_type, CheckpointType::ElementExists);
assert_eq!(elem_check.timeout_ms, Some(10000));
let js_check = Checkpoint::js_condition("document.readyState === 'complete'");
assert_eq!(js_check.checkpoint_type, CheckpointType::JsCondition);
}
#[test]
fn test_execution_plan() {
let plan = ExecutionPlan::new("Login to dashboard")
.add_step(PlannedStep::new(
"s1",
serde_json::json!({"Fill": {"selector": "input", "value": "user"}}),
))
.add_step(PlannedStep::new(
"s2",
serde_json::json!({"Click": "button"}),
))
.with_confidence(0.85)
.complete();
assert_eq!(plan.goal, "Login to dashboard");
assert_eq!(plan.step_count(), 2);
assert!(plan.is_complete);
assert_eq!(plan.confidence, 0.85);
}
#[test]
fn test_plan_execution_state() {
let mut state = PlanExecutionState::new();
state.complete_step("step1");
assert!(state.is_completed("step1"));
assert!(!state.is_failed("step1"));
assert_eq!(state.steps_executed(), 1);
state.fail_step("step2", "Selector not found");
assert!(state.is_failed("step2"));
assert!(state.has_failures());
assert_eq!(state.first_error(), Some("Selector not found"));
}
#[test]
fn test_checkpoint_result() {
let checkpoint = Checkpoint::url_contains("/success");
let passed = CheckpointResult::passed(checkpoint.clone(), 100);
assert!(passed.passed);
let failed =
CheckpointResult::failed(checkpoint, "URL did not match", 200).with_actual("/error");
assert!(!failed.passed);
assert_eq!(failed.actual_value, Some("/error".to_string()));
}
#[test]
fn test_plan_parsing() {
let json = serde_json::json!({
"goal": "Test automation",
"reasoning": "This plan will test the form",
"steps": [
{
"id": "s1",
"action": { "Click": "button" },
"confidence": 0.9,
"checkpoint": {
"type": "element_exists",
"expected": ".result"
}
}
],
"confidence": 0.85,
"is_complete": true
});
let plan = ExecutionPlan::from_json(&json).unwrap();
assert_eq!(plan.goal, "Test automation");
assert_eq!(plan.step_count(), 1);
assert!(plan.is_complete);
let step = &plan.steps[0];
assert_eq!(step.id, "s1");
assert!(step.checkpoint.is_some());
}
#[test]
fn test_replan_context() {
let context = ReplanContext::new(
"Complete checkout",
"click_payment",
"Payment button not found",
PageState::new("https://shop.com/cart", "Shopping Cart"),
)
.with_completed(vec!["add_to_cart".to_string()]);
let prompt = context.to_prompt();
assert!(prompt.contains("RE-PLANNING REQUIRED"));
assert!(prompt.contains("Complete checkout"));
assert!(prompt.contains("Payment button not found"));
assert!(prompt.contains("add_to_cart"));
}
#[test]
fn test_average_step_confidence() {
let plan = ExecutionPlan::new("Test")
.add_step(PlannedStep::new("s1", serde_json::json!({})).with_confidence(0.8))
.add_step(PlannedStep::new("s2", serde_json::json!({})).with_confidence(0.6));
assert_eq!(plan.average_step_confidence(), 0.7);
}
#[test]
fn test_critical_steps() {
let plan = ExecutionPlan::new("Test")
.add_step(PlannedStep::new("s1", serde_json::json!({})).with_critical(true))
.add_step(PlannedStep::new("s2", serde_json::json!({})).with_critical(false))
.add_step(PlannedStep::new("s3", serde_json::json!({})).with_critical(true));
let critical = plan.critical_steps();
assert_eq!(critical.len(), 2);
}
}