components/omnibox/composebox/composebox_query_controller.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 COMPONENTS_OMNIBOX_COMPOSEBOX_COMPOSEBOX_QUERY_CONTROLLER_H_
#define COMPONENTS_OMNIBOX_COMPOSEBOX_COMPOSEBOX_QUERY_CONTROLLER_H_

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

#include "base/memory/ref_counted_memory.h"
#include "base/memory/scoped_refptr.h"
#include "base/time/time.h"
#include "base/unguessable_token.h"
#include "components/endpoint_fetcher/endpoint_fetcher.h"
#include "components/lens/lens_overlay_mime_type.h"
#include "components/lens/lens_overlay_request_id_generator.h"
#include "services/network/public/cpp/shared_url_loader_factory.h"
#include "third_party/lens_server_proto/lens_overlay_client_context.pb.h"
#include "third_party/lens_server_proto/lens_overlay_cluster_info.pb.h"
#include "third_party/lens_server_proto/lens_overlay_server.pb.h"
#include "third_party/lens_server_proto/lens_overlay_surface.pb.h"

enum class SessionState {
  kNone = 0,
  kSessionStarted = 1,
  kSessionAbandoned = 2,
  kSubmittedQuery = 3,
};

enum class QueryControllerState {
  // The initial state, before NotifySessionStarted() is called.
  kOff = 0,
  // The cluster info request is in flight.
  kAwaitingClusterInfoResponse = 1,
  // The cluster info response has been received and is valid.
  kClusterInfoReceived = 2,
  // The cluster info response was not received, or the cluster info has
  // expired.
  kClusterInfoInvalid = 3,
};

// Upload status of a file.
enum class FileUploadStatus {
  // Not uploaded.
  kNotUploaded = 0,
  // File being processed.
  kProcessing = 1,
  // Failed validation - Terminal for this file attempt.
  kValidationFailed = 2,
  // Request sent to Lens server.
  kUploadStarted = 3,
  // Server confirmed successful receipt.
  kUploadSuccessful = 4,
  // Server or network error during upload - Terminal for this file attempt.
  kUploadFailed = 5,
};

// For upload error metrics.
enum class FileUploadErrorType {
  // Unknown.
  kUnknown = 0,
  // Browser error before/during request, not covered by validation.
  kBrowserProcessingError = 1,
  // Network-level issue (e.g., no connectivity, DNS failure).
  kNetworkError = 2,
  // Server returned an error (e.g., 5xx, specific API error).
  kServerError = 3,
  // Server rejected due to size after upload attempt - Considered terminal.
  kServerSizeLimitExceeded = 4,
  // Upload aborted by user deletion or session end.
  kAborted = 5,
};

namespace version_info {
enum class Channel;
}  // namespace version_info

namespace signin {
class IdentityManager;
}  // namespace signin

// Callback type alias for the OAuth headers created.
using OAuthHeadersCreatedCallback =
    base::OnceCallback<void(std::vector<std::string>)>;
// Callback type alias for the request body proto created.
using RequestBodyProtoCreatedCallback =
    base::OnceCallback<void(lens::LensOverlayServerRequest)>;
// Callback type alias for the upload progress.
using UploadProgressCallback =
    base::RepeatingCallback<void(uint64_t position, uint64_t total)>;
// Callback for when the query controller state changes.
using QueryControllerStateChangedCallback =
    base::RepeatingCallback<void(QueryControllerState state)>;
// Callback for when the file upload status changes.
using FileUploadStatusChangedCallback =
    base::RepeatingCallback<void(std::string file_token,
                                 FileUploadStatus status)>;

class ComposeboxQueryController {
 public:
  // Observer interface for the Page Handler to get updates on file upload
  class FileUploadStatusObserver : public base::CheckedObserver {
   public:
    virtual void OnFileUploadStatusChanged(
        const base::UnguessableToken& file_token,
        FileUploadStatus file_upload_status,
        const std::optional<FileUploadErrorType>& error_type) = 0;

   protected:
    ~FileUploadStatusObserver() override = default;
  };

  // Struct containing file information for a file upload.
  struct FileInfo {
   public:
    FileInfo();
    ~FileInfo();

    // Client-side unique identifier generated by UI. Used as the key in the
    // `active_files_` map.
    base::UnguessableToken file_token_;

    // The mime type of the file.
    lens::MimeType mime_type_;

    // Gets the file upload status.
    FileUploadStatus GetFileUploadStatus() const { return upload_status_; }

    // Gets the file upload error type.
    FileUploadErrorType GetFileUploadErrorType() const {
      return upload_error_type_;
    }

    // Gets the server response code.
    int GetResponseCode() const { return response_code_; }

    // Gets a pointer to the request ID for this request for testing.
    lens::LensOverlayRequestId* GetRequestIdForTesting() {
      return request_id_.get();
    }

   private:
    friend class ComposeboxQueryController;

    // Default to kNotUploaded, until UploadFile() is called.
    // Do not modify this field directly, use UpdateFileUploadStatus() instead.
    FileUploadStatus upload_status_ = FileUploadStatus::kNotUploaded;

    // The error type if the upload failed.
    FileUploadErrorType upload_error_type_ = FileUploadErrorType::kUnknown;

    // The request ID for this request. Set by StartFileUploadFlow().
    std::unique_ptr<lens::LensOverlayRequestId> request_id_;

    // When browser started the network request for the file upload.
    base::Time upload_network_request_start_time_;

    // When Lens server response was received.
    base::Time server_response_time_;

    // The network response code.
    int response_code_ = 0;

    // The request to be sent to the server. Will be set asynchronously after
    // StartFileUploadFlow() is called.
    std::unique_ptr<lens::LensOverlayServerRequest> request_body_;

    // The headers to attach to the request. Will be set asynchronously after
    // StartFileUploadFlow() is called.
    std::unique_ptr<std::vector<std::string>> request_headers_;

    // The access token fetcher used for getting OAuth for the file upload
    // request. Will be discarded after the OAuth headers are created.
    std::unique_ptr<signin::PrimaryAccountAccessTokenFetcher>
        file_upload_access_token_fetcher_;

    // The endpoint fetcher used for the file upload request.
    std::unique_ptr<endpoint_fetcher::EndpointFetcher>
        file_upload_endpoint_fetcher_;
  };

  ComposeboxQueryController(
      signin::IdentityManager* identity_manager,
      scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
      version_info::Channel channel);
  virtual ~ComposeboxQueryController();

  // Session management. Virtual for testing.
  virtual void NotifySessionStarted();
  virtual void NotifySessionAbandoned();

  // Observer management.
  void AddObserver(FileUploadStatusObserver* obs);
  void RemoveObserver(FileUploadStatusObserver* obs);

  SessionState session_state() { return session_state_; }

  // Triggers upload of the file with data and stores the file info in the
  // internal map. Call after setting the file info fields.
  void StartFileUploadFlow(std::unique_ptr<FileInfo> file_info,
                           scoped_refptr<base::RefCountedBytes> file_data);

 protected:
  // Returns the EndpointFetcher to use with the given params. Protected to
  // allow overriding in tests to mock server responses.
  virtual std::unique_ptr<endpoint_fetcher::EndpointFetcher>
  CreateEndpointFetcher(std::string request_string,
                        const GURL& fetch_url,
                        endpoint_fetcher::HttpMethod http_method,
                        base::TimeDelta timeout,
                        const std::vector<std::string>& request_headers,
                        const std::vector<std::string>& cors_exempt_headers,
                        UploadProgressCallback upload_progress_callback);

  // The internal state of the query controller. Protected to allow tests to
  // access the state. Do not modify this state directly, use
  // SetQueryControllerState() instead.
  QueryControllerState query_controller_state_ = QueryControllerState::kOff;

  // Callback for when the query controller state changes. Protected to allow
  // tests to set the callback.
  QueryControllerStateChangedCallback
      on_query_controller_state_changed_callback_;

  // The map of active files, keyed by the file token.
  // Protected to allow tests to access the files.
  std::map<base::UnguessableToken, std::unique_ptr<FileInfo>> active_files_;

 private:
  // Creates the client context for Lens requests.
  lens::LensOverlayClientContext CreateClientContext();

  // Fetches the OAuth headers and calls the callback with the headers. If the
  // OAuth cannot be retrieved (like if the user is not logged in), the callback
  // will be called with an empty vector. Returns the access token fetcher
  // making the request so it can be kept alive.
  std::unique_ptr<signin::PrimaryAccountAccessTokenFetcher>
  CreateOAuthHeadersAndContinue(OAuthHeadersCreatedCallback callback);

  // Gets an OAuth token for the cluster info request and proceeds with sending
  // a LensOverlayServerClusterInfoRequest to get the cluster info.
  void FetchClusterInfo();

  // Asynchronous handler for when the fetch cluster info request headers are
  // ready. Creates the endpoint fetcher and sends the cluster info request.
  void SendClusterInfoNetworkRequest(std::vector<std::string> request_headers);

  // Handles the response from the cluster info request.
  void HandleClusterInfoResponse(
      std::unique_ptr<endpoint_fetcher::EndpointResponse> response);

  // Sets the query controller state and notifies the query controller state
  // changed callback if it has changed.
  void SetQueryControllerState(QueryControllerState new_state);

  // Updates the file upload status and notifies the file upload status
  // observers with an optional error type if the upload failed.
  void UpdateFileUploadStatus(const base::UnguessableToken& file_token,
                              FileUploadStatus status,
                              std::optional<FileUploadErrorType> error_type);

  // Creates the request body proto and calls the callback with the request.
  void CreateFileUploadRequestBodyAndContinue(
      const base::UnguessableToken& file_token,
      scoped_refptr<base::RefCountedBytes> file_data,
      RequestBodyProtoCreatedCallback callback);

  // Asynchronous handler for when the file upload request body is ready.
  void OnUploadFileRequestBodyReady(const base::UnguessableToken& file_token,
                                    lens::LensOverlayServerRequest request);

  // Asynchronous handler for when the file upload request headers are ready.
  void OnUploadFileRequestHeadersReady(const base::UnguessableToken& file_token,
                                       std::vector<std::string> headers);

  // Sends the file upload request if the request body, headers, and cluster
  // info are ready.
  void MaybeSendFileUploadNetworkRequest(
      const base::UnguessableToken& file_token);

  // Creates the endpoint fetcher and sends the file upload network request.
  void SendFileUploadNetworkRequest(FileInfo* file_infon);

  // Handles the response from the file upload request.
  void HandleFileUploadResponse(
      const base::UnguessableToken& file_token,
      std::unique_ptr<endpoint_fetcher::EndpointResponse> response);

  // Return the file from `active_files_` map or nullptr if not found.
  FileInfo* GetFileInfo(const base::UnguessableToken& file_token);

  // The last received cluster info.
  std::optional<lens::LensOverlayClusterInfo> cluster_info_ = std::nullopt;

  // The endpoint fetcher used for the cluster info request.
  std::unique_ptr<endpoint_fetcher::EndpointFetcher>
      cluster_info_endpoint_fetcher_;

  // The access token fetcher used for getting OAuth for the cluster info
  // request. Will be discarded after the OAuth headers are created.
  std::unique_ptr<signin::PrimaryAccountAccessTokenFetcher>
      cluster_info_access_token_fetcher_;

  // Unowned IdentityManager for fetching access tokens. Could be null for
  // incognito profiles.
  const raw_ptr<signin::IdentityManager> identity_manager_;

  // TODO(420701010) Create SessionMetrics struct.
  base::Time session_start_time_;

  // The url loader factory to use for Lens network requests.
  scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory_;

  // The channel to use for Lens network requests.
  version_info::Channel channel_;

  // The request id generator for this query flow instance.
  lens::LensOverlayRequestIdGenerator request_id_generator_;

  // The session state.
  SessionState session_state_ = SessionState::kNone;

  // The observer list, managed via AddObserver() and RemoveObserver().
  base::ObserverList<FileUploadStatusObserver> observers_;

  // Task runner used to create the file upload request proto asynchronously.
  scoped_refptr<base::TaskRunner> create_request_task_runner_;

  base::WeakPtrFactory<ComposeboxQueryController> weak_ptr_factory_{this};
};

#endif  // COMPONENTS_OMNIBOX_COMPOSEBOX_COMPOSEBOX_QUERY_CONTROLLER_H_