From f14016302ecd75b9689b4a6fa9b9ab844ccb0dfc Mon Sep 17 00:00:00 2001 From: famez Date: Thu, 29 Jan 2026 22:20:50 +0100 Subject: [PATCH] Restored removed function (when clicking Esc to cancel current task, it throws an exception). --- pentestagent/agents/base_agent.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py index b752e2e..ea3c026 100644 --- a/pentestagent/agents/base_agent.py +++ b/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 \ No newline at end of file + 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) \ No newline at end of file