From e4165a45adbb0587f7cfeb8387521509890b95f4 Mon Sep 17 00:00:00 2001
From: Santiago Vendramini <sdvendramini@gmail.com>
Date: Fri, 6 Dec 2024 10:23:54 +0100
Subject: [PATCH] feat: CleanUpInProgressCommands moved inside the
 CommandsProcessingTask

With this change the function can be private.
---
 .../command_handler/include/command_handler.hpp | 17 +++++++++--------
 src/agent/src/agent.cpp                         | 12 ------------
 2 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/src/agent/command_handler/include/command_handler.hpp b/src/agent/command_handler/include/command_handler.hpp
index 090adc8d8c..f04d8456f1 100644
--- a/src/agent/command_handler/include/command_handler.hpp
+++ b/src/agent/command_handler/include/command_handler.hpp
@@ -49,6 +49,8 @@ namespace command_handler
             const auto executor = co_await boost::asio::this_coro::executor;
             std::unique_ptr<boost::asio::steady_timer> expTimer = std::make_unique<boost::asio::steady_timer>(executor);
 
+            CleanUpInProgressCommands(ReportCommandResult);
+
             while (m_keepRunning.load())
             {
                 auto cmd = GetCommandFromQueue();
@@ -85,8 +87,12 @@ namespace command_handler
             }
         }
 
-        void CleanUpInProgressCommands(
-            std::function<void(const std::optional<std::vector<module_command::CommandEntry>>&)> callback)
+        /// @brief Stops the command handler
+        void Stop();
+
+    private:
+        template<typename T>
+        void CleanUpInProgressCommands(std::function<void(T&)> ReportCommandResult)
         {
             auto cmds = m_commandStore.GetCommandByStatus(module_command::Status::IN_PROGRESS);
 
@@ -97,16 +103,11 @@ namespace command_handler
                     cmd.ExecutionResult.ErrorCode = module_command::Status::FAILURE;
                     cmd.ExecutionResult.Message = "Agent stopped during execution";
                     m_commandStore.UpdateCommand(cmd);
+                    ReportCommandResult(cmd);
                 }
             }
-
-            callback(cmds);
         }
 
-        /// @brief Stops the command handler
-        void Stop();
-
-    private:
         /// @brief Indicates whether the command handler is running or not
         std::atomic<bool> m_keepRunning = true;
 
diff --git a/src/agent/src/agent.cpp b/src/agent/src/agent.cpp
index 254713aff6..1a6a440fc1 100644
--- a/src/agent/src/agent.cpp
+++ b/src/agent/src/agent.cpp
@@ -60,18 +60,6 @@ Agent::Agent(const std::string& configFilePath, std::unique_ptr<ISignalHandler>
 
     m_taskManager.Start(
         m_configurationParser->GetConfig<size_t>("agent", "thread_count").value_or(config::DEFAULT_THREAD_COUNT));
-
-    m_commandHandler.CleanUpInProgressCommands(
-        [this](const auto& inProgressCommands)
-        {
-            if (inProgressCommands != std::nullopt)
-            {
-                for (auto& cmd : *inProgressCommands)
-                {
-                    ReportCommandResult(cmd, m_messageQueue);
-                }
-            }
-        });
 }
 
 Agent::~Agent()