mirror of
https://github.com/GH05TCREW/pentestagent.git
synced 2026-03-08 06:44:11 +00:00
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
"""Agent state management for GhostCrew."""
|
|
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
class AgentState(Enum):
|
|
"""Possible states for an agent."""
|
|
|
|
IDLE = "idle"
|
|
THINKING = "thinking"
|
|
EXECUTING = "executing"
|
|
WAITING_INPUT = "waiting_input"
|
|
COMPLETE = "complete"
|
|
ERROR = "error"
|
|
|
|
|
|
@dataclass
|
|
class StateTransition:
|
|
"""Represents a state transition."""
|
|
|
|
from_state: AgentState
|
|
to_state: AgentState
|
|
timestamp: datetime = field(default_factory=datetime.now)
|
|
reason: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class AgentStateManager:
|
|
"""Manages agent state and transitions."""
|
|
|
|
current_state: AgentState = AgentState.IDLE
|
|
history: List[StateTransition] = field(default_factory=list)
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Valid state transitions
|
|
VALID_TRANSITIONS = {
|
|
AgentState.IDLE: [AgentState.THINKING, AgentState.ERROR],
|
|
AgentState.THINKING: [
|
|
AgentState.EXECUTING,
|
|
AgentState.WAITING_INPUT,
|
|
AgentState.COMPLETE,
|
|
AgentState.ERROR,
|
|
],
|
|
AgentState.EXECUTING: [AgentState.THINKING, AgentState.ERROR],
|
|
AgentState.WAITING_INPUT: [
|
|
AgentState.THINKING,
|
|
AgentState.COMPLETE,
|
|
AgentState.ERROR,
|
|
],
|
|
AgentState.COMPLETE: [AgentState.IDLE],
|
|
AgentState.ERROR: [AgentState.IDLE],
|
|
}
|
|
|
|
def can_transition_to(self, new_state: AgentState) -> bool:
|
|
"""Check if a transition to the new state is valid."""
|
|
valid_next_states = self.VALID_TRANSITIONS.get(self.current_state, [])
|
|
return new_state in valid_next_states
|
|
|
|
def transition_to(
|
|
self, new_state: AgentState, reason: Optional[str] = None
|
|
) -> bool:
|
|
"""
|
|
Transition to a new state.
|
|
|
|
Args:
|
|
new_state: The state to transition to
|
|
reason: Optional reason for the transition
|
|
|
|
Returns:
|
|
True if transition was successful, False otherwise
|
|
"""
|
|
if not self.can_transition_to(new_state):
|
|
return False
|
|
|
|
transition = StateTransition(
|
|
from_state=self.current_state, to_state=new_state, reason=reason
|
|
)
|
|
self.history.append(transition)
|
|
self.current_state = new_state
|
|
return True
|
|
|
|
def force_transition(self, new_state: AgentState, reason: Optional[str] = None):
|
|
"""Force a transition regardless of validity (use with caution)."""
|
|
transition = StateTransition(
|
|
from_state=self.current_state,
|
|
to_state=new_state,
|
|
reason=f"FORCED: {reason}" if reason else "FORCED",
|
|
)
|
|
self.history.append(transition)
|
|
self.current_state = new_state
|
|
|
|
def reset(self):
|
|
"""Reset state to IDLE."""
|
|
self.current_state = AgentState.IDLE
|
|
self.history.clear()
|
|
self.metadata.clear()
|
|
|
|
def is_terminal(self) -> bool:
|
|
"""Check if current state is a terminal state."""
|
|
return self.current_state in [AgentState.COMPLETE, AgentState.ERROR]
|
|
|
|
def is_active(self) -> bool:
|
|
"""Check if agent is actively processing."""
|
|
return self.current_state in [AgentState.THINKING, AgentState.EXECUTING]
|
|
|
|
def get_state_duration(self) -> float:
|
|
"""Get duration in current state (seconds)."""
|
|
if not self.history:
|
|
return 0.0
|
|
|
|
last_transition = self.history[-1]
|
|
return (datetime.now() - last_transition.timestamp).total_seconds()
|