diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py
index caed3bc..990defb 100644
--- a/pentestagent/interface/tui.py
+++ b/pentestagent/interface/tui.py
@@ -3186,6 +3186,7 @@ Be concise. Use the actual data from notes."""
         self._set_status("thinking", "assist")
 
         try:
+            last_tool_message: ToolMessage
             async for response in self.agent.assist(message):
                 if self._should_stop:
                     self._add_system("[!] Stopped by user")
@@ -3207,7 +3208,7 @@ Be concise. Use the actual data from notes."""
                         if call.name == "finish":
                             continue  # Skip - summary shown as final message
                         args_str = str(call.arguments)
-                        self._add_tool(call.name, args_str)
+                        last_tool_message = self._add_tool(call.name, args_str)
 
                 # Show tool results (displayed after execution completes)
                 # Skip 'finish' tool - its result is shown as the final summary
@@ -3216,11 +3217,11 @@ Be concise. Use the actual data from notes."""
                         if result.tool_name == "finish":
                             continue  # Skip - summary shown separately
                         if result.success:
-                            self._add_tool_result(
+                            self._add_tool_result(last_tool_message, 
                                 result.tool_name, result.result or "Done"
                             )
                         else:
-                            self._add_tool_result(
+                            self._add_tool_result(last_tool_message, 
                                 result.tool_name, f"Error: {result.error}"
                             )