From e0074f1f47b61f7e876c7046dc66a940821929ae Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Thu, 20 Jun 2024 11:24:37 +0200
Subject: [PATCH] refact: unified serialization logic

---
 src/agent/mod.rs               | 26 +++++++++++++++++++++--
 src/agent/serialization/mod.rs | 38 ----------------------------------
 src/agent/serialization/xml.rs |  3 +--
 src/agent/state/history.rs     |  6 ++++--
 src/agent/state/mod.rs         |  3 ++-
 5 files changed, 31 insertions(+), 45 deletions(-)

diff --git a/src/agent/mod.rs b/src/agent/mod.rs
index cc2f959..e84513f 100644
--- a/src/agent/mod.rs
+++ b/src/agent/mod.rs
@@ -1,8 +1,9 @@
+use std::collections::HashMap;
+
 use anyhow::Result;
 use colored::Colorize;
 
 use model::{Client, Options};
-use serialization::Invocation;
 use state::State;
 use task::Task;
 
@@ -12,6 +13,27 @@ mod serialization;
 pub mod state;
 pub mod task;
 
+#[derive(Debug, Default, Clone, PartialEq)]
+pub struct Invocation {
+    pub action: String,
+    pub attributes: Option<HashMap<String, String>>,
+    pub payload: Option<String>,
+}
+
+impl Invocation {
+    pub fn new(
+        action: String,
+        attributes: Option<HashMap<String, String>>,
+        payload: Option<String>,
+    ) -> Self {
+        Self {
+            action,
+            attributes,
+            payload,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct AgentOptions {
     pub max_iterations: usize,
@@ -123,7 +145,7 @@ impl Agent {
         for inv in invocations {
             // avoid running the same command twince in a row
             if let Some(p) = prev.as_ref() {
-                if inv.is_same(p) {
+                if inv.eq(p) {
                     println!(".");
                     continue;
                 }
diff --git a/src/agent/serialization/mod.rs b/src/agent/serialization/mod.rs
index e3d45cd..62fc266 100644
--- a/src/agent/serialization/mod.rs
+++ b/src/agent/serialization/mod.rs
@@ -1,39 +1 @@
-use std::collections::HashMap;
-
-use xml::serialize_invocation;
-
 pub(crate) mod xml;
-
-#[derive(Debug, Default, Clone)]
-pub struct Invocation {
-    pub action: String,
-    pub attributes: Option<HashMap<String, String>>,
-    pub payload: Option<String>,
-
-    serialized: String,
-}
-
-impl Invocation {
-    pub fn new(
-        action: String,
-        attributes: Option<HashMap<String, String>>,
-        payload: Option<String>,
-    ) -> Self {
-        let mut zelf = Self {
-            action,
-            attributes,
-            payload,
-            serialized: "".to_string(),
-        };
-        zelf.serialized = serialize_invocation(&zelf);
-        zelf
-    }
-
-    pub fn as_serialized_str(&self) -> &str {
-        return self.serialized.as_str();
-    }
-
-    pub fn is_same(&self, other: &Invocation) -> bool {
-        self.serialized == other.serialized
-    }
-}
diff --git a/src/agent/serialization/xml.rs b/src/agent/serialization/xml.rs
index d2b19db..55ead11 100644
--- a/src/agent/serialization/xml.rs
+++ b/src/agent/serialization/xml.rs
@@ -7,10 +7,9 @@ use regex::Regex;
 use crate::agent::{
     namespaces::Action,
     state::storage::{Storage, StorageType, CURRENT_TAG, PREVIOUS_TAG},
+    Invocation,
 };
 
-use super::Invocation;
-
 lazy_static! {
     pub static ref XML_ATTRIBUTES_PARSER: Regex = Regex::new(r#"(?m)(([^=]+)="([^"]+)")"#).unwrap();
 }
diff --git a/src/agent/state/history.rs b/src/agent/state/history.rs
index 1bfaad9..31d5e28 100644
--- a/src/agent/state/history.rs
+++ b/src/agent/state/history.rs
@@ -1,6 +1,6 @@
 use anyhow::Result;
 
-use crate::agent::{model::Message, serialization::Invocation};
+use crate::agent::{model::Message, Invocation};
 
 #[derive(Debug, Clone, Default)]
 pub struct Execution {
@@ -47,7 +47,9 @@ impl Execution {
         if let Some(response) = self.response.as_ref() {
             messages.push(Message::Agent(response.to_string()));
         } else if let Some(invocation) = self.invocation.as_ref() {
-            messages.push(Message::Agent(invocation.as_serialized_str().to_string()));
+            messages.push(Message::Agent(
+                crate::agent::serialization::xml::serialize_invocation(invocation),
+            ));
         }
 
         messages.push(Message::Feedback(if let Some(err) = &self.error {
diff --git a/src/agent/state/mod.rs b/src/agent/state/mod.rs
index 07f8f38..acc9820 100644
--- a/src/agent/state/mod.rs
+++ b/src/agent/state/mod.rs
@@ -12,8 +12,9 @@ use colored::Colorize;
 use super::{
     model::Message,
     namespaces::{self, Namespace},
-    serialization::{self, Invocation},
+    serialization,
     task::Task,
+    Invocation,
 };
 use history::{Execution, History};
 use storage::Storage;