chrome/browser/permissions/prediction_service/permissions_aiv1_handler.cc - chromium/src.git - Git at Google

 // 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.

 #include "chrome/browser/permissions/prediction_service/permissions_aiv1_handler.h"

 #include "base/check_is_test.h"
 #include "base/containers/fixed_flat_set.h"
 #include "base/metrics/histogram_functions.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/sequence_checker.h"
 #include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/profiles/profile.h"
 #include "components/optimization_guide/proto/features/permissions_ai.pb.h"
 #include "components/permissions/request_type.h"

 namespace permissions {

 namespace {
 using ::optimization_guide::ModelBasedCapabilityKey;
 using ::optimization_guide::SessionConfigParams;
 using ::optimization_guide::proto::PermissionsAiRequest;
 using ::optimization_guide::proto::PermissionsAiResponse;
 using ::optimization_guide::proto::PermissionType;
 using EligibilityReason = ::optimization_guide::OnDeviceModelEligibilityReason;

 constexpr ModelBasedCapabilityKey kFeatureKey =
     ModelBasedCapabilityKey::kPermissionsAi;

 constexpr SessionConfigParams kSessionConfigParams = SessionConfigParams{
     .execution_mode = SessionConfigParams::ExecutionMode::kOnDeviceOnly,
 };

 // Max delay for permissions AI model execution; inquiries that need more time
 // get cancelled.
 constexpr base::TimeDelta kMaxModelExecutionDuration = base::Seconds(5);

 void LogOnDeviceModelPreviousSessionFinishedInTime(bool success) {
   base::UmaHistogramBoolean("Permissions.AIv1.PreviousSessionFinishedInTime",
                             success);
 }

 void LogOnDeviceModelSessionCreationSuccess(bool session_created) {
   base::UmaHistogramBoolean("Permissions.AIv1.SessionCreationSuccess",
                             session_created);
 }

 void LogOnDeviceModelExecutionSuccessAndTime(
     bool success,
     base::TimeTicks session_execution_start_time) {
   base::UmaHistogramBoolean("Permissions.AIv1.ExecutionSuccess", success);
   if (success) {
     base::UmaHistogramMediumTimes(
         "Permissions.AIv1.ExecutionDuration",
         base::TimeTicks::Now() - session_execution_start_time);
   }
 }

 void LogOnDeviceModelExecutionParse(bool success) {
   base::UmaHistogramBoolean("Permissions.AIv1.ResponseParseSuccess", success);
 }

 void LogOnDeviceModelAvailabilityAtInquiryTime(bool model_available) {
   base::UmaHistogramBoolean("Permissions.AIv1.AvailableAtInquiryTime",
                             model_available);
 }

 void LogOnDeviceModelExecutionTimedOut(bool timed_out) {
   base::UmaHistogramBoolean("Permissions.AIv1.ExecutionTimedOut", timed_out);
 }

 PermissionType GetPermissionType(permissions::RequestType request_type) {
   switch (request_type) {
     case permissions::RequestType::kNotifications:
       return PermissionType::PERMISSION_TYPE_NOTIFICATIONS;
     case permissions::RequestType::kGeolocation:
       return PermissionType::PERMISSION_TYPE_GEOLOCATION;
     default:
       return PermissionType::PERMISSION_TYPE_NOT_SPECIFIED;
   }
 }

 bool IsOnDeviceModelAvailable(EligibilityReason reason) {
   return reason == EligibilityReason::kSuccess;
 }
 }  // namespace

 // Manages sessions and related data. This will gift us more flexibility, i.e.
 // it can allow us to have multiple sessions at some point in the future at the
 // same time and also to easily cancel session without fearing asynchronous
 // calls to OnModelExecutionComplete meddling with a potentially new permission
 // request.
 class PermissionsAiv1Handler::EvaluationTask {
  public:
   explicit EvaluationTask(OptimizationGuideKeyedService* optimization_guide) {
     DETACH_FROM_SEQUENCE(sequence_checker_);
     session_ =
         optimization_guide->StartSession(kFeatureKey, kSessionConfigParams);
   }

   ~EvaluationTask() {
     DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
     // Explicitly invalidate weak pointers to prevent callbacks that may be
     // triggered by the destructor logic.
     weak_ptr_factory_.InvalidateWeakPtrs();
     if (inquire_on_device_model_callback_) {
       std::move(inquire_on_device_model_callback_).Run(std::nullopt);
     }
   }

   bool IsActive() { return session_ != nullptr; }

   void ExecuteModel(
       std::string rendered_text,
       permissions::RequestType request_type,
       base::OnceCallback<void(std::optional<PermissionsAiResponse>)> callback,
       base::OneShotTimer* execution_timer) {
     DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
     if (!session_) {
       std::move(callback).Run(std::nullopt);
       return;
     }

     inquire_on_device_model_callback_ = std::move(callback);

     PermissionsAiRequest request;
     request.set_rendered_text(std::move(rendered_text));
     request.set_permission_type(GetPermissionType(request_type));

     session_execution_start_time_ = base::TimeTicks::Now();
     session_->ExecuteModel(
         request,
         base::BindRepeating(
             &PermissionsAiv1Handler::EvaluationTask::OnModelExecutionComplete,
             weak_ptr_factory_.GetWeakPtr(), execution_timer));
   }

  private:
   void OnModelExecutionComplete(
       base::OneShotTimer* execution_timer,
       optimization_guide::OptimizationGuideModelStreamingExecutionResult
           result) {
     DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

     // This is a non-error response, but since the result is not fully available
     // yet, we defer callback execution until we get a complete response.
     if (result.response.has_value() && !result.response->is_complete) {
       return;
     }

     if (execution_timer) {
       execution_timer->Stop();
       LogOnDeviceModelExecutionTimedOut(/*timed_out=*/false);
     }
     // Since we do not want to reuse a once used session, lets make sure that we
     // destroy it when the model execution has finished.
     session_.reset();

     // If there is no callback available anymore, we do not need to do any work
     // here.
     if (!inquire_on_device_model_callback_) {
       return;
     }

     LogOnDeviceModelExecutionSuccessAndTime(/*success=*/
                                             result.response.has_value(),
                                             session_execution_start_time_);
     if (!result.response.has_value()) {
       VLOG(1)
           << "[PermissionsAIv1] OnModelExecutionComplete failed with error: "
           << static_cast<int>(result.response.error().error());
       std::move(inquire_on_device_model_callback_).Run(std::nullopt);
       return;
     }

     std::optional<PermissionsAiResponse> permissions_ai_response =
         optimization_guide::ParsedAnyMetadata<PermissionsAiResponse>(
             result.response->response);
     LogOnDeviceModelExecutionParse(
         /*success=*/permissions_ai_response.has_value());

     if (!permissions_ai_response.has_value()) {
       VLOG(1) << "[PermissionsAIv1] OnModelExecutionComplete failed while "
                  "parsing the response proto.";
       std::move(inquire_on_device_model_callback_).Run(std::nullopt);
       return;
     }

     std::move(inquire_on_device_model_callback_).Run(permissions_ai_response);
   }

   SEQUENCE_CHECKER(sequence_checker_);

   std::unique_ptr<optimization_guide::OptimizationGuideModelExecutor::Session>
       session_;
   base::TimeTicks session_execution_start_time_;
   base::OnceCallback<void(
       std::optional<optimization_guide::proto::PermissionsAiResponse>)>
       inquire_on_device_model_callback_;
   base::WeakPtrFactory<PermissionsAiv1Handler::EvaluationTask>
       weak_ptr_factory_{this};
 };

 PermissionsAiv1Handler::PermissionsAiv1Handler(
     OptimizationGuideKeyedService* optimization_guide)
     : optimization_guide_(optimization_guide),
       execution_timer_(std::make_unique<base::OneShotTimer>()) {}

 PermissionsAiv1Handler::~PermissionsAiv1Handler() {
   execution_timer_->Stop();
 }

 bool PermissionsAiv1Handler::IsModelExecutionInProgress() {
   return evaluation_task_ && evaluation_task_->IsActive();
 }

 void PermissionsAiv1Handler::InquireAiOnDeviceModel(
     std::string rendered_text,
     permissions::RequestType request_type,
     base::OnceCallback<void(std::optional<PermissionsAiResponse>)> callback) {
   if (!optimization_guide_) {
     // If optimization_guide_ is a nullptr then we cannot do anything here at
     // all.
     std::move(callback).Run(std::nullopt);
     return;
   }

   EligibilityReason model_state =
       optimization_guide_->GetOnDeviceModelEligibility(kFeatureKey);
   bool model_available_at_inquiry_time = IsOnDeviceModelAvailable(model_state);
   LogOnDeviceModelAvailabilityAtInquiryTime(model_available_at_inquiry_time);

   if (model_available_at_inquiry_time) {
     if (IsModelExecutionInProgress()) {
       LogOnDeviceModelPreviousSessionFinishedInTime(/*success=*/false);
       // TODO(crbug.com/382447738): It can happen that a new inquiry comes
       // before the previous finishes its execution. To avoid unexpected
       // behavior return `std::nullopt` which means another type of prediction_service logic
       // will be executed.
       std::move(callback).Run(std::nullopt);
       return;
     }
     LogOnDeviceModelPreviousSessionFinishedInTime(/*success=*/true);
   }

   // We make sure by recreating the session on every inquiry that no
   // callback executions or timing data linked to the previous request gets
   // accidentally mixed up with the new request when we cancel lengthy model
   // executions.
   evaluation_task_ = std::make_unique<EvaluationTask>(optimization_guide_);

   LogOnDeviceModelSessionCreationSuccess(
       /*session_created=*/evaluation_task_->IsActive());

   if (evaluation_task_->IsActive()) {
     evaluation_task_->ExecuteModel(std::move(rendered_text), request_type,
                                    std::move(callback), execution_timer_.get());

     // We check again to make sure the callback did not immediately return.
     if (evaluation_task_->IsActive()) {
       execution_timer_->Start(
           FROM_HERE, kMaxModelExecutionDuration,
           base::BindOnce(&PermissionsAiv1Handler::CancelModelExecution,
                          weak_ptr_factory_.GetWeakPtr()));
     }
     return;
   }

   // We end up here if the model download has not started/ended yet or
   // there is a transient error. If we did not download the model yet, trying
   // to create the session will have started the download already so we do
   // nothing here and just return. We also return for the other two cases,
   // hoping the next inquiry will have better luck.
   std::move(callback).Run(std::nullopt);
 }

 void PermissionsAiv1Handler::CancelModelExecution() {
   evaluation_task_.reset();
   execution_timer_->Stop();
   LogOnDeviceModelExecutionTimedOut(/*timed_out=*/true);
 }

 void PermissionsAiv1Handler::set_execution_timer_for_testing(
     std::unique_ptr<base::OneShotTimer> execution_timer) {
   CHECK_IS_TEST();
   execution_timer_ = std::move(execution_timer);
 }

 }  // namespace permissions
	// 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.

	#include "chrome/browser/permissions/prediction_service/permissions_aiv1_handler.h"

	#include "base/check_is_test.h"
	#include "base/containers/fixed_flat_set.h"
	#include "base/metrics/histogram_functions.h"
	#include "base/metrics/histogram_macros.h"
	#include "base/sequence_checker.h"
	#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
	#include "chrome/browser/profiles/profile.h"
	#include "components/optimization_guide/proto/features/permissions_ai.pb.h"
	#include "components/permissions/request_type.h"

	namespace permissions {

	namespace {
	using ::optimization_guide::ModelBasedCapabilityKey;
	using ::optimization_guide::SessionConfigParams;
	using ::optimization_guide::proto::PermissionsAiRequest;
	using ::optimization_guide::proto::PermissionsAiResponse;
	using ::optimization_guide::proto::PermissionType;
	using EligibilityReason = ::optimization_guide::OnDeviceModelEligibilityReason;

	constexpr ModelBasedCapabilityKey kFeatureKey =
	ModelBasedCapabilityKey::kPermissionsAi;

	constexpr SessionConfigParams kSessionConfigParams = SessionConfigParams{
	.execution_mode = SessionConfigParams::ExecutionMode::kOnDeviceOnly,
	};

	// Max delay for permissions AI model execution; inquiries that need more time
	// get cancelled.
	constexpr base::TimeDelta kMaxModelExecutionDuration = base::Seconds(5);

	void LogOnDeviceModelPreviousSessionFinishedInTime(bool success) {
	base::UmaHistogramBoolean("Permissions.AIv1.PreviousSessionFinishedInTime",
	success);
	}

	void LogOnDeviceModelSessionCreationSuccess(bool session_created) {
	base::UmaHistogramBoolean("Permissions.AIv1.SessionCreationSuccess",
	session_created);
	}

	void LogOnDeviceModelExecutionSuccessAndTime(
	bool success,
	base::TimeTicks session_execution_start_time) {
	base::UmaHistogramBoolean("Permissions.AIv1.ExecutionSuccess", success);
	if (success) {
	base::UmaHistogramMediumTimes(
	"Permissions.AIv1.ExecutionDuration",
	base::TimeTicks::Now() - session_execution_start_time);
	}
	}

	void LogOnDeviceModelExecutionParse(bool success) {
	base::UmaHistogramBoolean("Permissions.AIv1.ResponseParseSuccess", success);
	}

	void LogOnDeviceModelAvailabilityAtInquiryTime(bool model_available) {
	base::UmaHistogramBoolean("Permissions.AIv1.AvailableAtInquiryTime",
	model_available);
	}

	void LogOnDeviceModelExecutionTimedOut(bool timed_out) {
	base::UmaHistogramBoolean("Permissions.AIv1.ExecutionTimedOut", timed_out);
	}

	PermissionType GetPermissionType(permissions::RequestType request_type) {
	switch (request_type) {
	case permissions::RequestType::kNotifications:
	return PermissionType::PERMISSION_TYPE_NOTIFICATIONS;
	case permissions::RequestType::kGeolocation:
	return PermissionType::PERMISSION_TYPE_GEOLOCATION;
	default:
	return PermissionType::PERMISSION_TYPE_NOT_SPECIFIED;
	}
	}

	bool IsOnDeviceModelAvailable(EligibilityReason reason) {
	return reason == EligibilityReason::kSuccess;
	}
	} // namespace

	// Manages sessions and related data. This will gift us more flexibility, i.e.
	// it can allow us to have multiple sessions at some point in the future at the
	// same time and also to easily cancel session without fearing asynchronous
	// calls to OnModelExecutionComplete meddling with a potentially new permission
	// request.
	class PermissionsAiv1Handler::EvaluationTask {
	public:
	explicit EvaluationTask(OptimizationGuideKeyedService* optimization_guide) {
	DETACH_FROM_SEQUENCE(sequence_checker_);
	session_ =
	optimization_guide->StartSession(kFeatureKey, kSessionConfigParams);
	}

	~EvaluationTask() {
	DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
	// Explicitly invalidate weak pointers to prevent callbacks that may be
	// triggered by the destructor logic.
	weak_ptr_factory_.InvalidateWeakPtrs();
	if (inquire_on_device_model_callback_) {
	std::move(inquire_on_device_model_callback_).Run(std::nullopt);
	}
	}

	bool IsActive() { return session_ != nullptr; }

	void ExecuteModel(
	std::string rendered_text,
	permissions::RequestType request_type,
	base::OnceCallback<void(std::optional<PermissionsAiResponse>)> callback,
	base::OneShotTimer* execution_timer) {
	DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
	if (!session_) {
	std::move(callback).Run(std::nullopt);
	return;
	}

	inquire_on_device_model_callback_ = std::move(callback);

	PermissionsAiRequest request;
	request.set_rendered_text(std::move(rendered_text));
	request.set_permission_type(GetPermissionType(request_type));

	session_execution_start_time_ = base::TimeTicks::Now();
	session_->ExecuteModel(
	request,
	base::BindRepeating(
	&PermissionsAiv1Handler::EvaluationTask::OnModelExecutionComplete,
	weak_ptr_factory_.GetWeakPtr(), execution_timer));
	}

	private:
	void OnModelExecutionComplete(
	base::OneShotTimer* execution_timer,
	optimization_guide::OptimizationGuideModelStreamingExecutionResult
	result) {
	DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

	// This is a non-error response, but since the result is not fully available
	// yet, we defer callback execution until we get a complete response.
	if (result.response.has_value() && !result.response->is_complete) {
	return;
	}

	if (execution_timer) {
	execution_timer->Stop();
	LogOnDeviceModelExecutionTimedOut(/timed_out=/false);
	}
	// Since we do not want to reuse a once used session, lets make sure that we
	// destroy it when the model execution has finished.
	session_.reset();

	// If there is no callback available anymore, we do not need to do any work
	// here.
	if (!inquire_on_device_model_callback_) {
	return;
	}

	LogOnDeviceModelExecutionSuccessAndTime(/success=/
	result.response.has_value(),
	session_execution_start_time_);
	if (!result.response.has_value()) {
	VLOG(1)
	<< "[PermissionsAIv1] OnModelExecutionComplete failed with error: "
	<< static_cast<int>(result.response.error().error());
	std::move(inquire_on_device_model_callback_).Run(std::nullopt);
	return;
	}

	std::optional<PermissionsAiResponse> permissions_ai_response =
	optimization_guide::ParsedAnyMetadata<PermissionsAiResponse>(
	result.response->response);
	LogOnDeviceModelExecutionParse(
	/success=/permissions_ai_response.has_value());

	if (!permissions_ai_response.has_value()) {
	VLOG(1) << "[PermissionsAIv1] OnModelExecutionComplete failed while "
	"parsing the response proto.";
	std::move(inquire_on_device_model_callback_).Run(std::nullopt);
	return;
	}

	std::move(inquire_on_device_model_callback_).Run(permissions_ai_response);
	}

	SEQUENCE_CHECKER(sequence_checker_);

	std::unique_ptr<optimization_guide::OptimizationGuideModelExecutor::Session>
	session_;
	base::TimeTicks session_execution_start_time_;
	base::OnceCallback<void(
	std::optional<optimization_guide::proto::PermissionsAiResponse>)>
	inquire_on_device_model_callback_;
	base::WeakPtrFactory<PermissionsAiv1Handler::EvaluationTask>
	weak_ptr_factory_{this};
	};

	PermissionsAiv1Handler::PermissionsAiv1Handler(
	OptimizationGuideKeyedService* optimization_guide)
	: optimization_guide_(optimization_guide),
	execution_timer_(std::make_unique<base::OneShotTimer>()) {}

	PermissionsAiv1Handler::~PermissionsAiv1Handler() {
	execution_timer_->Stop();
	}

	bool PermissionsAiv1Handler::IsModelExecutionInProgress() {
	return evaluation_task_ && evaluation_task_->IsActive();
	}

	void PermissionsAiv1Handler::InquireAiOnDeviceModel(
	std::string rendered_text,
	permissions::RequestType request_type,
	base::OnceCallback<void(std::optional<PermissionsAiResponse>)> callback) {
	if (!optimization_guide_) {
	// If optimization_guide_ is a nullptr then we cannot do anything here at
	// all.
	std::move(callback).Run(std::nullopt);
	return;
	}

	EligibilityReason model_state =
	optimization_guide_->GetOnDeviceModelEligibility(kFeatureKey);
	bool model_available_at_inquiry_time = IsOnDeviceModelAvailable(model_state);
	LogOnDeviceModelAvailabilityAtInquiryTime(model_available_at_inquiry_time);

	if (model_available_at_inquiry_time) {
	if (IsModelExecutionInProgress()) {
	LogOnDeviceModelPreviousSessionFinishedInTime(/success=/false);
	// TODO(crbug.com/382447738): It can happen that a new inquiry comes
	// before the previous finishes its execution. To avoid unexpected
	// behavior return `std::nullopt` which means another type of prediction_service logic
	// will be executed.
	std::move(callback).Run(std::nullopt);
	return;
	}
	LogOnDeviceModelPreviousSessionFinishedInTime(/success=/true);
	}

	// We make sure by recreating the session on every inquiry that no
	// callback executions or timing data linked to the previous request gets
	// accidentally mixed up with the new request when we cancel lengthy model
	// executions.
	evaluation_task_ = std::make_unique<EvaluationTask>(optimization_guide_);

	LogOnDeviceModelSessionCreationSuccess(
	/session_created=/evaluation_task_->IsActive());

	if (evaluation_task_->IsActive()) {
	evaluation_task_->ExecuteModel(std::move(rendered_text), request_type,
	std::move(callback), execution_timer_.get());

	// We check again to make sure the callback did not immediately return.
	if (evaluation_task_->IsActive()) {
	execution_timer_->Start(
	FROM_HERE, kMaxModelExecutionDuration,
	base::BindOnce(&PermissionsAiv1Handler::CancelModelExecution,
	weak_ptr_factory_.GetWeakPtr()));
	}
	return;
	}

	// We end up here if the model download has not started/ended yet or
	// there is a transient error. If we did not download the model yet, trying
	// to create the session will have started the download already so we do
	// nothing here and just return. We also return for the other two cases,
	// hoping the next inquiry will have better luck.
	std::move(callback).Run(std::nullopt);
	}

	void PermissionsAiv1Handler::CancelModelExecution() {
	evaluation_task_.reset();
	execution_timer_->Stop();
	LogOnDeviceModelExecutionTimedOut(/timed_out=/true);
	}

	void PermissionsAiv1Handler::set_execution_timer_for_testing(
	std::unique_ptr<base::OneShotTimer> execution_timer) {
	CHECK_IS_TEST();
	execution_timer_ = std::move(execution_timer);
	}

	} // namespace permissions