From e1b25c9357b40e9b35ac3d2b2c0f27436898484e Mon Sep 17 00:00:00 2001
From: GH05TCREW <harmasic@gmail.com>
Date: Sun, 28 Dec 2025 04:47:21 -0700
Subject: [PATCH] refactor(config): unify agent and orchestrator iteration env
 vars

---
 .env.example                             | 8 +++++++-
 pentestagent/agents/base_agent.py        | 4 ++--
 pentestagent/agents/crew/orchestrator.py | 4 ++--
 pentestagent/agents/crew/worker_pool.py  | 4 ++--
 pentestagent/config/constants.py         | 6 ++++--
 pentestagent/config/settings.py          | 4 ++--
 6 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/.env.example b/.env.example
index e1722cc..3e69fc5 100644
--- a/.env.example
+++ b/.env.example
@@ -20,4 +20,10 @@ PENTESTAGENT_MODEL=gpt-5
 # PENTESTAGENT_EMBEDDINGS=local
 
 # Settings
-PENTESTAGENT_DEBUG=false
\ No newline at end of file
+PENTESTAGENT_DEBUG=false
+
+# Agent max iterations (regular agent + crew workers, default: 30)
+# PENTESTAGENT_AGENT_MAX_ITERATIONS=30
+
+# Orchestrator max iterations (crew mode coordinator, default: 50)
+# PENTESTAGENT_ORCHESTRATOR_MAX_ITERATIONS=50
\ No newline at end of file
diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py
index 32d1b1d..8e165da 100644
--- a/pentestagent/agents/base_agent.py
+++ b/pentestagent/agents/base_agent.py
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, AsyncIterator, List, Optional
 
-from ..config.constants import DEFAULT_MAX_ITERATIONS
+from ..config.constants import AGENT_MAX_ITERATIONS
 from .state import AgentState, AgentStateManager
 
 if TYPE_CHECKING:
@@ -78,7 +78,7 @@ class BaseAgent(ABC):
         llm: "LLM",
         tools: List["Tool"],
         runtime: "Runtime",
-        max_iterations: int = DEFAULT_MAX_ITERATIONS,
+        max_iterations: int = AGENT_MAX_ITERATIONS,
     ):
         """
         Initialize the base agent.
diff --git a/pentestagent/agents/crew/orchestrator.py b/pentestagent/agents/crew/orchestrator.py
index 7c164fb..a62f11b 100644
--- a/pentestagent/agents/crew/orchestrator.py
+++ b/pentestagent/agents/crew/orchestrator.py
@@ -3,7 +3,7 @@
 import json
 from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional
 
-from ...config.constants import DEFAULT_MAX_ITERATIONS
+from ...config.constants import ORCHESTRATOR_MAX_ITERATIONS
 from ...knowledge.graph import ShadowGraph
 from ..prompts import pa_crew
 from .models import CrewState, WorkerCallback
@@ -156,7 +156,7 @@ class CrewOrchestrator:
         final_report = ""
 
         try:
-            while iteration < DEFAULT_MAX_ITERATIONS:
+            while iteration < ORCHESTRATOR_MAX_ITERATIONS:
                 iteration += 1
 
                 response = await self.llm.generate(
diff --git a/pentestagent/agents/crew/worker_pool.py b/pentestagent/agents/crew/worker_pool.py
index 059b624..4b40529 100644
--- a/pentestagent/agents/crew/worker_pool.py
+++ b/pentestagent/agents/crew/worker_pool.py
@@ -109,7 +109,7 @@ class WorkerPool:
         worker_runtime = LocalRuntime()
         await worker_runtime.start()
 
-        from ...config.constants import WORKER_MAX_ITERATIONS
+        from ...config.constants import AGENT_MAX_ITERATIONS
 
         agent = PentestAgentAgent(
             llm=self.llm,
@@ -117,7 +117,7 @@ class WorkerPool:
             runtime=worker_runtime,  # Use isolated runtime
             target=self.target,
             rag_engine=self.rag_engine,
-            max_iterations=WORKER_MAX_ITERATIONS,
+            max_iterations=AGENT_MAX_ITERATIONS,
         )
 
         try:
diff --git a/pentestagent/config/constants.py b/pentestagent/config/constants.py
index c05d31a..e1402d8 100644
--- a/pentestagent/config/constants.py
+++ b/pentestagent/config/constants.py
@@ -56,8 +56,10 @@ DEFAULT_TEMPERATURE = 0.7
 DEFAULT_MAX_TOKENS = 4096
 
 # Agent Defaults
-DEFAULT_MAX_ITERATIONS = int(os.environ.get("PENTESTAGENT_MAX_ITERATIONS", "50"))
-WORKER_MAX_ITERATIONS = int(os.environ.get("PENTESTAGENT_WORKER_MAX_ITERATIONS", "10"))
+AGENT_MAX_ITERATIONS = int(os.environ.get("PENTESTAGENT_AGENT_MAX_ITERATIONS", "30"))
+ORCHESTRATOR_MAX_ITERATIONS = int(
+    os.environ.get("PENTESTAGENT_ORCHESTRATOR_MAX_ITERATIONS", "50")
+)
 
 # File Extensions
 KNOWLEDGE_TEXT_EXTENSIONS = [".txt", ".md"]
diff --git a/pentestagent/config/settings.py b/pentestagent/config/settings.py
index 4c6a38b..1058188 100644
--- a/pentestagent/config/settings.py
+++ b/pentestagent/config/settings.py
@@ -6,7 +6,7 @@ from pathlib import Path
 from typing import List, Optional
 
 from .constants import (
-    DEFAULT_MAX_ITERATIONS,
+    AGENT_MAX_ITERATIONS,
     DEFAULT_MAX_TOKENS,
     DEFAULT_MODEL,
     DEFAULT_TEMPERATURE,
@@ -40,7 +40,7 @@ class Settings:
     docker_image: str = "ghcr.io/gh05tcrew/pentestagent:kali"
 
     # Agent Settings
-    max_iterations: int = DEFAULT_MAX_ITERATIONS
+    max_iterations: int = AGENT_MAX_ITERATIONS
 
     # VPN Settings
     vpn_config_path: Optional[Path] = None