Merge pull request #23 from famez/main

Added removed function

Dockerfile

@@ -46,9 +46,6 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
-# Copy application code
-COPY . .
-
 # Create non-root user for security
 RUN useradd -m -s /bin/bash pentestagent && \
     chown -R pentestagent:pentestagent /app
@@ -60,6 +57,8 @@ USER pentestagent
 
 RUN playwright install
 
+# Copy application code
+COPY . .
 
 # Expose any needed ports
 EXPOSE 8080

pentestagent/agents/base_agent.py

@@ -940,4 +940,31 @@ Call create_plan with the new steps OR feasible=False."""
         return "\n".join(parts)
 
     def get_state(self) -> AgentState:
-        return self.state_manager.current_state
+        return self.state_manager.current_state
+
+    def cleanup_after_cancel(self) -> None:
+        """
+        Clean up agent state after a cancellation.
+
+        Removes the cancelled request and any pending tool calls from
+        conversation history to prevent stale responses from contaminating
+        the next conversation.
+        """
+        # Remove incomplete messages from the end of conversation
+        while self.conversation_history:
+            last_msg = self.conversation_history[-1]
+            # Remove assistant message with tool calls (incomplete tool execution)
+            if last_msg.role == "assistant" and last_msg.tool_calls:
+                self.conversation_history.pop()
+            # Remove orphaned tool_result messages
+            elif last_msg.role == "tool":
+                self.conversation_history.pop()
+            # Remove the user message that triggered the cancelled request
+            elif last_msg.role == "user":
+                self.conversation_history.pop()
+                break  # Stop after removing the user message
+            else:
+                break
+
+        # Reset state to idle
+        self.state_manager.transition_to(AgentState.IDLE)