Replace eval() of model output in task handlers with literal parser (close #1567)

superagi/agent/output_handler.py

@@ -145,8 +145,7 @@ def __init__(self, agent_execution_id: int, agent_config: dict):
         self.agent_config = agent_config
 
     def handle(self, session, assistant_reply):
-        assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply)
-        tasks = eval(assistant_reply)
+        tasks = JsonCleaner.parse_array_section(assistant_reply)
         tasks = np.array(tasks).flatten().tolist()
         for task in reversed(tasks):
             self.task_queue.add_task(task)
@@ -176,8 +175,7 @@ def __init__(self, agent_execution_id: int, agent_config: dict):
         self.agent_config = agent_config
 
     def handle(self, session, assistant_reply):
-        assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply)
-        tasks = eval(assistant_reply)
+        tasks = JsonCleaner.parse_array_section(assistant_reply)
         self.task_queue.clear_tasks()
         for task in reversed(tasks):
             self.task_queue.add_task(task)

superagi/agent/queue_step_handler.py

@@ -74,9 +74,9 @@ def _consume_from_queue(self, task_queue: TaskQueue):
             task_queue.complete_task("PROCESSED")
 
     def _process_reply(self, task_queue: TaskQueue, assistant_reply: str):
-        assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply)
-        print("Queue reply:", assistant_reply)
-        task_array = np.array(eval(assistant_reply)).flatten().tolist()
+        tasks = JsonCleaner.parse_array_section(assistant_reply)
+        print("Queue reply:", tasks)
+        task_array = np.array(tasks).flatten().tolist()
         for task in task_array:
             task_queue.add_task(str(task))
             logger.info("RAMRAM: Added task to queue: ", task)

superagi/helper/json_cleaner.py

@@ -1,3 +1,4 @@
+import ast
 import json
 import re
 from superagi.lib.logger import logger
@@ -7,6 +8,39 @@
 
 class JsonCleaner:
 
+    @classmethod
+    def parse_array_section(cls, input_str: str = ""):
+        """
+        Safely parse the array section of an LLM reply into a Python list.
+
+        This is the safe replacement for ``eval()`` on model output. Model
+        replies are only ever expected to contain a literal array of task
+        strings (e.g. ``["task1", "task2"]`` or, because OpenAI emits
+        ``str(dict)``-style payloads, single-quoted ``['task1', 'task2']``).
+        ``ast.literal_eval`` accepts both quoting styles but, unlike ``eval``,
+        evaluates *only* Python literals, so it can never execute function
+        calls, imports or other code embedded in the reply
+        (e.g. ``[__import__('os').system('id')]`` raises instead of running).
+
+        Args:
+            input_str (str): The (already array-sliced) model reply.
+
+        Returns:
+            list: The parsed list of tasks.
+
+        Raises:
+            ValueError: If the input is not a literal Python/JSON list.
+        """
+        array_section = cls.extract_json_array_section(input_str)
+        try:
+            parsed = ast.literal_eval(array_section)
+        except (ValueError, SyntaxError, TypeError, MemoryError, RecursionError):
+            logger.error("JsonCleaner.parse_array_section: unable to parse array section safely")
+            raise ValueError("Model reply did not contain a valid literal array")
+        if not isinstance(parsed, (list, tuple)):
+            raise ValueError("Model reply did not contain a valid literal array")
+        return list(parsed)
+
     @classmethod
     def clean_boolean(cls, input_str: str = ""):
         """

tests/unit_tests/agent/test_output_handler.py

@@ -125,6 +125,30 @@ def test_task_output_handle_method(extract_json_array_section_mock, get_tasks_mo
     assert response.status == "PENDING"
 
 
+# Security regression: TaskOutputHandler must not execute code embedded in the
+# model reply. A prompt-injected/adversarial model can return a string that is a
+# valid Python *expression* (e.g. a call to os.system); the handler must reject
+# it as a non-literal instead of evaluating it.
+@patch.object(TaskQueue, 'add_task')
+@patch.object(TaskQueue, 'get_tasks')
+def test_task_output_handle_rejects_code_execution(get_tasks_mock, add_task_mock):
+    agent_execution_id = 1
+    agent_config = {"agent_id": 2}
+    # Payload that, under eval(), would invoke os.system and run a command.
+    malicious_reply = "[__import__('os').system('echo pwned')]"
+    get_tasks_mock.return_value = []
+    handler = TaskOutputHandler(agent_execution_id, agent_config)
+    session_mock = MagicMock()
+
+    with patch('os.system') as os_system_mock:
+        with pytest.raises(ValueError):
+            handler.handle(session_mock, malicious_reply)
+        os_system_mock.assert_not_called()
+
+    # No tasks should have been queued from the malicious payload.
+    add_task_mock.assert_not_called()
+
+
 # Test for ReplaceTaskOutputHandler
 @patch.object(TaskQueue, 'clear_tasks')
 @patch.object(TaskQueue, 'add_task')
@@ -152,3 +176,28 @@ def test_handle_method(extract_json_array_section_mock, get_tasks_mock, add_task
     assert add_task_mock.call_count == len(tasks)
     get_tasks_mock.assert_called_once()
     assert response.status == "PENDING"
+
+
+# Security regression: ReplaceTaskOutputHandler must not execute code embedded in
+# the model reply. A prompt-injected/adversarial model can return a string that is
+# a valid Python *expression* (e.g. a call to subprocess.getoutput); the handler
+# must reject it as a non-literal instead of evaluating it.
+@patch.object(TaskQueue, 'clear_tasks')
+@patch.object(TaskQueue, 'add_task')
+@patch.object(TaskQueue, 'get_tasks')
+def test_replace_task_output_handle_rejects_code_execution(get_tasks_mock, add_task_mock, clear_tasks_mock):
+    agent_execution_id = 1
+    agent_config = {}
+    # Payload that, under eval(), would invoke subprocess.getoutput and run a command.
+    malicious_reply = "[__import__('subprocess').getoutput('echo pwned')]"
+    get_tasks_mock.return_value = []
+    handler = ReplaceTaskOutputHandler(agent_execution_id, agent_config)
+    session_mock = MagicMock()
+
+    with patch('subprocess.getoutput') as getoutput_mock:
+        with pytest.raises(ValueError):
+            handler.handle(session_mock, malicious_reply)
+        getoutput_mock.assert_not_called()
+
+    # No tasks should have been queued from the malicious payload.
+    add_task_mock.assert_not_called()

tests/unit_tests/agent/test_queue_step_handler.py

@@ -60,6 +60,32 @@ def test_add_to_queue(task_queue_mock, agent_execution_feed_mock, queue_step_han
     queue_step_handler._process_reply.assert_called_once_with(task_queue_mock, '{"reply": ["task1", "task2"]}')
 
 
+# Security regression: _process_reply must not execute code embedded in the model
+# reply. A prompt-injected/adversarial model can return a string that is a valid
+# Python *expression* (e.g. a call to os.system); the queue step must reject it as a
+# non-literal instead of evaluating it.
+def test_process_reply_rejects_code_execution(queue_step_handler):
+    task_queue_mock = Mock()
+    # Payload that, under eval(), would invoke os.system and run a command.
+    malicious_reply = "[__import__('os').system('echo pwned')]"
+
+    with patch('os.system') as os_system_mock:
+        with pytest.raises(ValueError):
+            queue_step_handler._process_reply(task_queue_mock, malicious_reply)
+        os_system_mock.assert_not_called()
+
+    # No task should have been added from the malicious payload.
+    task_queue_mock.add_task.assert_not_called()
+
+
+def test_process_reply_parses_literal_array(queue_step_handler):
+    # A legitimate model reply (a literal list of task strings) must still be parsed
+    # and enqueued, including the single-quoted str(dict)-style form OpenAI emits.
+    task_queue_mock = Mock()
+    queue_step_handler._process_reply(task_queue_mock, "['task1', 'task2']")
+    assert task_queue_mock.add_task.call_count == 2
+
+
 @patch("superagi.agent.queue_step_handler.TaskQueue")
 @patch("superagi.agent.queue_step_handler.AgentExecutionFeed")
 def test_consume_from_queue(task_queue_mock, agent_execution_feed_mock, queue_step_handler, step_tool):