Fix Agent Loop Bug

Author: yashab-cyber

hackbot/core/pdf_report.py

@@ -1103,7 +1103,11 @@ def _normalize_tool_history(tool_history: Optional[List[Dict[str, Any]]]) -> Lis
             cmd = str(entry.get("command", "") or "").strip()
             tool = str(entry.get("tool", "") or "").strip()
             if not tool and cmd:
-                tool = cmd.split()[0]
+                parts = cmd.split()
+                if parts[0] == "sudo" and len(parts) > 1:
+                    tool = parts[1]
+                else:
+                    tool = parts[0]
 
             out = dict(entry)
             out["tool"] = tool or "unknown"

tests/test_modes.py

@@ -3,7 +3,11 @@
 import pytest
 
 from hackbot.modes.plan import PlanMode, PLAN_TEMPLATES
-from hackbot.modes.agent import Severity, Finding
+from hackbot.modes.agent import AgentMode, Severity, Finding
+from hackbot.core.engine import AIEngine
+from hackbot.config import HackBotConfig, AgentConfig, AIConfig
+from hackbot.core.runner import ToolResult
+from unittest.mock import MagicMock, patch
 
 
 def test_plan_templates_exist():
@@ -49,3 +53,95 @@ def test_finding_creation():
     d = finding.to_dict()
     assert d["severity"] == "High"
     assert d["title"] == "SQL Injection"
+
+
+@pytest.fixture
+def mock_agent():
+    """Create an AgentMode instance with mocked dependencies for testing loop logic."""
+    config = HackBotConfig()
+    config.ai = AIConfig(provider="mock", model="mock")
+    config.agent = AgentConfig(allowed_tools=["echo", "nmap"], safe_mode=False)
+    
+    engine = MagicMock(spec=AIEngine)
+    engine.chat.return_value = "Mocked AI response"
+    
+    with patch("hackbot.core.vulndb.VulnDB"), \
+         patch("hackbot.core.cve.CVELookup"):
+        agent = AgentMode(engine=engine, config=config)
+        agent.target = "127.0.0.1"
+        agent.conversation = MagicMock()
+        agent.is_running = True
+        return agent
+
+def test_process_actions_loop_skips_same_round_duplicates(mock_agent):
+    """Test that identical commands within one round are skipped."""
+    mock_agent._execute_action = MagicMock(return_value=("output", ToolResult(
+        tool="echo", command="echo test", stdout="test", stderr="", 
+        return_code=0, duration=0.1, success=True
+    )))
+    
+    actions = [
+        {"action": "execute", "command": "echo test"},
+        {"action": "execute", "command": "echo test"},
+        {"action": "execute", "command": "echo test"},
+    ]
+    
+    # We only care about the single-round processing here, so mock Chat to stop the loop
+    mock_agent.engine.chat.return_value = '{"action": "complete"}'
+    
+    mock_agent._process_actions_loop(actions, max_rounds=2)
+    
+    # execute_action should only be called once, the others are deduplicated in same round
+    assert mock_agent._execute_action.call_count == 1
+
+def test_process_actions_loop_skips_session_duplicates(mock_agent):
+    """Test that commands executed >=2 times in a session are skipped."""
+    mock_agent._command_history = {"echo test": 2} # Already run twice
+    
+    mock_agent._execute_action = MagicMock()
+    mock_agent.engine.chat.return_value = '{"action": "complete"}'
+    
+    actions = [
+        {"action": "execute", "command": "echo test"},
+    ]
+    mock_agent._process_actions_loop(actions, max_rounds=2)
+    
+    # Should be skipped entirely
+    mock_agent._execute_action.assert_not_called()
+
+def test_command_history_increments(mock_agent):
+    """Test that executing an action increments command history."""
+    # Ensure runner is a mock that won't actually execute
+    mock_agent.runner = MagicMock()
+    mock_agent.runner.execute.return_value = ToolResult(
+        tool="echo", command="echo hello", stdout="hello", stderr="", 
+        return_code=0, duration=0.1, success=True
+    )
+    
+    action = {"action": "execute", "command": "echo hello"}
+    
+    assert "echo hello" not in mock_agent._command_history
+    mock_agent._execute_action(action)
+    assert mock_agent._command_history["echo hello"] == 1
+    
+    mock_agent._execute_action(action)
+    assert mock_agent._command_history["echo hello"] == 2
+
+def test_should_nudge_capped_per_step(mock_agent):
+    """Test that _should_nudge only returns True once per step."""
+    response_with_tool = "I will run nmap now."
+    
+    # First check should be True
+    assert mock_agent._should_nudge(response_with_tool) is True
+    # Second check should be False (already nudged)
+    assert mock_agent._should_nudge(response_with_tool) is False
+    
+    # Reset count (as happens at top of step())
+    mock_agent._nudge_count = 0
+    assert mock_agent._should_nudge(response_with_tool) is True
+
+def test_should_nudge_ignores_non_tools(mock_agent):
+    """Test that _should_nudge returns False if no tool names are present."""
+    response_no_tool = "I have begun assessing the target based on the scope provided."
+    # Should be false since it doesn't mention nmap/sqlmap/etc
+    assert mock_agent._should_nudge(response_no_tool) is False