diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py
index e9172b5..b752e2e 100644
--- a/pentestagent/agents/base_agent.py
+++ b/pentestagent/agents/base_agent.py
@@ -938,3 +938,6 @@ Call create_plan with the new steps OR feasible=False."""
             else:
                 parts.append(f"Error: {r.error}")
         return "\n".join(parts)
+
+    def get_state(self) -> AgentState:
+        return self.state_manager.current_state
\ No newline at end of file
diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py
index df8941c..2b00b67 100644
--- a/pentestagent/interface/tui.py
+++ b/pentestagent/interface/tui.py
@@ -2954,6 +2954,9 @@ Be concise. Use the actual data from notes."""
         self._set_status("thinking", "agent")
 
         try:
+
+            from ..agents.base_agent import AgentState
+
             async for response in self.agent.agent_loop(task):
                 if self._should_stop:
                     self._add_system("[!] Stopped by user")
@@ -2999,11 +3002,11 @@ Be concise. Use the actual data from notes."""
                             )
 
                 # Check state
-                if self.agent.state.value == "waiting_input":
+                if self.agent.get_state() == AgentState.WAITING_INPUT:
                     self._set_status("waiting")
                     self._add_system("? Awaiting input...")
                     break
-                elif self.agent.state.value == "complete":
+                elif self.agent.get_state() == AgentState.COMPLETE:
                     break
 
                 self._set_status("thinking")