chrome/browser/contextual_cueing/zero_state_suggestions_page_data.h

// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef CHROME_BROWSER_CONTEXTUAL_CUEING_ZERO_STATE_SUGGESTIONS_PAGE_DATA_H_
#define CHROME_BROWSER_CONTEXTUAL_CUEING_ZERO_STATE_SUGGESTIONS_PAGE_DATA_H_

#include <optional>
#include <string>
#include <vector>

#include "base/callback_list.h"
#include "base/time/time.h"
#include "components/optimization_guide/content/browser/page_content_proto_provider.h"
#include "components/optimization_guide/core/optimization_guide_decision.h"
#include "components/optimization_guide/core/optimization_guide_model_executor.h"
#include "components/optimization_guide/core/optimization_metadata.h"
#include "content/public/browser/page_user_data.h"

class OptimizationGuideKeyedService;

namespace content_extraction {
struct InnerTextResult;
}  // namespace content_extraction

namespace optimization_guide {
class ModelQualityLogEntry;
}  // namespace optimization_guide

namespace optimization_guide::proto {
class ZeroStateSuggestionsRequest;
}  // namespace optimization_guide::proto

namespace contextual_cueing {

using GlicSuggestionsCallbackList =
    base::OnceCallbackList<void(std::optional<std::vector<std::string>>)>;
using GlicSuggestionsCallback = GlicSuggestionsCallbackList::CallbackType;

// Processes zero state suggestions for GLIC, scoped to the given page.
class ZeroStateSuggestionsPageData
    : public content::PageUserData<ZeroStateSuggestionsPageData> {
 public:
  ZeroStateSuggestionsPageData(const ZeroStateSuggestionsPageData&) = delete;
  ZeroStateSuggestionsPageData& operator=(const ZeroStateSuggestionsPageData&) =
      delete;
  ~ZeroStateSuggestionsPageData() override;

  // Explicitly fetch suggestions for this page.
  void FetchSuggestions(bool is_fre, GlicSuggestionsCallback callback);

 private:
  friend class content::PageUserData<ZeroStateSuggestionsPageData>;
  friend class ZeroStateSuggestionsPageDataTest;

  // Note that this constructor initiates extracting page content.
  explicit ZeroStateSuggestionsPageData(content::Page& page);

  // Called when inner text is extracted.
  void OnReceivedInnerText(
      std::unique_ptr<content_extraction::InnerTextResult> result);

  // Called when annotated page content is received.
  void OnReceivedAnnotatedPageContent(
      std::optional<optimization_guide::AIPageContentResult> content);

  // Called when on-demand metadata is received.
  void OnReceivedOptimizationMetadata(
      const GURL& url,
      const base::flat_map<
          optimization_guide::proto::OptimizationType,
          optimization_guide::OptimizationGuideDecisionWithMetadata>&
          decisions);

  // Send out suggestions request, if all necessary fetches are complete.
  void RequestSuggestionsIfComplete();

  // Return suggestions result by invoking `suggestions_callbacks_`, if all
  // necessary fetches are complete.
  void ProcessSuggestionsIfComplete();

  // If `optimization_metadata_` contains everything necessary to determine a
  // suggestions result, run `suggestions_callbacks_` to return those
  // suggestions. This method itself also returns true if suggestions are sent
  // via the callbacks as a result of execution.
  bool ReturnSuggestionsFromOptimizationMetadataIfPossible();

  // Called when a zero state suggestions server response is received.
  void OnModelExecutionResponse(
      optimization_guide::OptimizationGuideModelExecutionResult result,
      std::unique_ptr<optimization_guide::ModelQualityLogEntry> log_entry);

  // Tracks the status of inner text and annotated page content fetches, which
  // are needed in `RequestSuggestionsIfComplete()`.
  bool inner_text_done_ = false;
  std::unique_ptr<content_extraction::InnerTextResult> inner_text_result_;
  bool annotated_page_content_done_ = false;
  std::optional<optimization_guide::AIPageContentResult>
      annotated_page_content_;

  // Tracks the status of fetches for on-demand optimization metadata, and
  // suggestions from model execution service, which are needed for
  // `ProcessSuggestionsIfComplete()`.
  bool optimization_metadata_done_ = false;
  optimization_guide::OptimizationGuideDecision optimization_decision_;
  optimization_guide::OptimizationMetadata optimization_metadata_;
  bool mes_suggestions_done_ = false;
  std::unique_ptr<optimization_guide::OptimizationGuideModelExecutionResult>
      mes_suggestions_result_;

  // Tracks the state for a request.
  base::TimeTicks begin_time_;
  std::optional<optimization_guide::proto::ZeroStateSuggestionsRequest>
      suggestions_request_;
  GlicSuggestionsCallbackList suggestions_callbacks_;

  std::optional<std::vector<std::string>> cached_suggestions_;

  // Not owned and guaranteed to outlive `this`.
  raw_ptr<OptimizationGuideKeyedService> optimization_guide_keyed_service_ =
      nullptr;
  base::WeakPtrFactory<ZeroStateSuggestionsPageData> weak_ptr_factory_{this};

  PAGE_USER_DATA_KEY_DECL();
};

}  // namespace contextual_cueing

#endif  // CHROME_BROWSER_CONTEXTUAL_CUEING_ZERO_STATE_SUGGESTIONS_PAGE_DATA_H_