Fix leak in AITextSession::Responder

Fixed: 356364157
Change-Id: I7051c1aebc3a6f829b2a200b283ba7276d8aa99a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5751304
Commit-Queue: Tsuyoshi Horo <horo@chromium.org>
Reviewed-by: Jiacheng Guo <gjc@google.com>
Reviewed-by: Mingyu Lei <leimy@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1335816}

third_party/blink/renderer/modules/ai/ai_text_session.cc

@@ -23,6 +23,7 @@
 #include "third_party/blink/renderer/modules/ai/ai_metrics.h"
 #include "third_party/blink/renderer/modules/ai/exception_helpers.h"
 #include "third_party/blink/renderer/platform/bindings/exception_state.h"
+#include "third_party/blink/renderer/platform/context_lifecycle_observer.h"
 #include "third_party/blink/renderer/platform/heap/garbage_collected.h"
 #include "third_party/blink/renderer/platform/heap/self_keep_alive.h"
 #include "third_party/blink/renderer/platform/mojo/heap_mojo_receiver.h"
@@ -37,15 +38,19 @@ using mojom::blink::ModelStreamingResponseStatus;
 // result through a promise.
 class AITextSession::Responder final
     : public GarbageCollected<AITextSession::Responder>,
-      public blink::mojom::blink::ModelStreamingResponder {
+      public blink::mojom::blink::ModelStreamingResponder,
+      public ContextLifecycleObserver {
  public:
   explicit Responder(ScriptState* script_state)
       : resolver_(MakeGarbageCollected<ScriptPromiseResolver<IDLString>>(
             script_state)),
-        receiver_(this, ExecutionContext::From(script_state)) {}
+        receiver_(this, ExecutionContext::From(script_state)) {
+    SetContextLifecycleNotifier(ExecutionContext::From(script_state));
+  }
   ~Responder() override = default;
 
-  void Trace(Visitor* visitor) const {
+  void Trace(Visitor* visitor) const override {
+    ContextLifecycleObserver::Trace(visitor);
     visitor->Trace(resolver_);
     visitor->Trace(receiver_);
   }
@@ -85,14 +90,23 @@ class AITextSession::Responder final
           AIMetrics::GetAISessionResponseCallbackCountMetricName(
               AIMetrics::AISessionType::kText),
           response_callback_count_);
-      keep_alive_.Clear();
+      Cleanup();
       return;
     }
     // When the status is kOngoing, update the response with the latest value.
     response_ = text;
   }
 
+  // ContextLifecycleObserver implementation.
+  void ContextDestroyed() override { Cleanup(); }
+
  private:
+  void Cleanup() {
+    resolver_ = nullptr;
+    receiver_.reset();
+    keep_alive_.Clear();
+  }
+
   Member<ScriptPromiseResolver<IDLString>> resolver_;
   WTF::String response_;
   int response_callback_count_;

third_party/blink/web_tests/LeakExpectations

@@ -46,9 +46,6 @@ crbug.com/1385278 [ Linux ] external/wpt/resource-timing/iframe-failed-commit.ht
 # Sheriff 2023-04-04
 crbug.com/1430360 [ Linux ] external/wpt/custom-elements/form-associated/label-delegatesFocus.html [ Failure Pass ]
 
-# Gardener 2024-07-30
-crbug.com/356364157 [ Linux ] wpt_internal/ai/destroy-api.any.html [ Failure ]
-
 ###########################################################################
 # WARNING: Memory leaks must be fixed asap. Sheriff is expected to revert #
 # culprit CLs instead of suppressing the leaks. If you have any question, #

third_party/blink/web_tests/wpt_internal/ai/prompt-api-context-destroyed.window.js

@@ -0,0 +1,12 @@
+promise_test(async t => {
+  // Make sure the prompt api is enabled.
+  assert_true(!!ai);
+  // Create the iframe and append it to the document.
+  const iframe = document.createElement('iframe');
+  document.childNodes[document.childNodes.length - 1].appendChild(iframe);
+
+  const session = await iframe.contentWindow.ai.createTextSession();
+  session.prompt('hello');
+  // Detach the iframe.
+  iframe.remove();
+}, 'Detaching iframe while runing prompt() should not cause memory leak');