[Autofill Assistant] Failing |GetAutofillAgent| calls

There are some crashes from |WebController| when resetting the keyboard
state through the |AutofillAgent|. These can happen if the
|RenderFrameHost| is no longer valid when the final callback to clean up
the state is executed.

Bug: 1119918
Change-Id: I9bcd69ae125a1ab7918bc5671cb3b95c4cc82d71
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2396121
Commit-Queue: Sandro Maggi <sandromaggi@google.com>
Reviewed-by: Marian Fechete <marianfe@google.com>
Reviewed-by: Maxim Kolosovskiy  <kolos@chromium.org>
Cr-Commit-Position: refs/heads/master@{#805275}

components/autofill_assistant/browser/web/web_controller.cc

@@ -1600,8 +1600,10 @@ void WebController::OnInternalWaitForDocumentToBecomeInteractive(
 
 WebController::ScopedAssistantActionStateRunning::
     ScopedAssistantActionStateRunning(
-        autofill::ContentAutofillDriver* content_autofill_driver)
-    : content_autofill_driver_(content_autofill_driver) {
+        content::WebContents* web_contents,
+        content::RenderFrameHost* render_frame_host)
+    : content::WebContentsObserver(web_contents),
+      render_frame_host_(render_frame_host) {
   SetAssistantActionState(/* running= */ true);
 }
 
@@ -1610,13 +1612,21 @@ WebController::ScopedAssistantActionStateRunning::
   SetAssistantActionState(/* running= */ false);
 }
 
+void WebController::ScopedAssistantActionStateRunning::RenderFrameDeleted(
+    content::RenderFrameHost* render_frame_host) {
+  if (render_frame_host_ == render_frame_host)
+    render_frame_host_ = nullptr;
+}
+
 void WebController::ScopedAssistantActionStateRunning::SetAssistantActionState(
     bool running) {
-  if (content_autofill_driver_) {
-    // TODO(b/153625351): We assume the |ContentAutofillDriver| is still valid
-    // at this point. This assumes the |RenderFrameHost| does not get destroyed
-    // during the execution of a web action (e.g. clicking an element).
-    content_autofill_driver_->GetAutofillAgent()->SetAssistantActionState(
+  if (render_frame_host_ == nullptr)
+    return;
+
+  ContentAutofillDriver* content_autofill_driver =
+      ContentAutofillDriver::GetForRenderFrameHost(render_frame_host_);
+  if (content_autofill_driver != nullptr) {
+    content_autofill_driver->GetAutofillAgent()->SetAssistantActionState(
         running);
   }
 }
@@ -1645,7 +1655,7 @@ WebController::GetAssistantActionRunningStateRetainingCallback(
 
   auto scoped_assistant_action_state_running =
       std::make_unique<ScopedAssistantActionStateRunning>(
-          content_autofill_driver);
+          web_contents_, element_result.container_frame_host);
 
   return base::BindOnce(
       &WebController::RetainAssistantActionRunningStateAndExecuteCallback,

components/autofill_assistant/browser/web/web_controller.h

@@ -27,13 +27,13 @@
 #include "components/autofill_assistant/browser/web/element_position_getter.h"
 #include "components/autofill_assistant/browser/web/element_rect_getter.h"
 #include "components/autofill_assistant/browser/web/web_controller_worker.h"
+#include "content/public/browser/web_contents_observer.h"
 #include "third_party/icu/source/common/unicode/umachine.h"
 #include "url/gurl.h"
 
 namespace autofill {
 class AutofillProfile;
 class CreditCard;
-class ContentAutofillDriver;
 struct FormData;
 struct FormFieldData;
 }  // namespace autofill
@@ -254,11 +254,13 @@ class WebController {
 
   // RAII object that sets the action state to "running" when the object is
   // allocated and to "not running" when it gets deallocated.
-  class ScopedAssistantActionStateRunning {
+  class ScopedAssistantActionStateRunning
+      : private content::WebContentsObserver {
    public:
     explicit ScopedAssistantActionStateRunning(
-        autofill::ContentAutofillDriver* content_autofill_driver);
-    ~ScopedAssistantActionStateRunning();
+        content::WebContents* web_contents,
+        content::RenderFrameHost* render_frame_host);
+    ~ScopedAssistantActionStateRunning() override;
 
     ScopedAssistantActionStateRunning(
         const ScopedAssistantActionStateRunning&) = delete;
@@ -268,7 +270,11 @@ class WebController {
    private:
     void SetAssistantActionState(bool running);
 
-    autofill::ContentAutofillDriver* content_autofill_driver_;
+    // Overrides content::WebContentsObserver:
+    void RenderFrameDeleted(
+        content::RenderFrameHost* render_frame_host) override;
+
+    content::RenderFrameHost* render_frame_host_;
   };
 
   void OnJavaScriptResult(