Skip to content

Commit a87dce0

Browse files
dslee414Chromium LUCI CQ
authored andcommitted
[M138] Built-in API: propagate QuotaExceededError info
This info will be used for QuotaExceededError.requested and QuotaExceededError.quota in a follow-up CL. Spec: whatwg/webidl#1465 (cherry picked from commit 644db35) Bug: 416303646, 406162261 Fixed: 420952090 Change-Id: Id1d414722d2f53de0a566e8b453acdd97f501714 Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6568411 Reviewed-by: Mike Wasserman <[email protected]> Commit-Queue: Daseul Lee <[email protected]> Reviewed-by: Emily Stark <[email protected]> Reviewed-by: Ayu Ishii <[email protected]> Cr-Original-Commit-Position: refs/heads/main@{#1466096} Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6597962 Reviewed-by: Joe Mason <[email protected]> Auto-Submit: Daseul Lee <[email protected]> Commit-Queue: Rubber Stamper <[email protected]> Bot-Commit: Rubber Stamper <[email protected]> Cr-Commit-Position: refs/branch-heads/7204@{#416} Cr-Branched-From: d5de512-refs/heads/main@{#1465706}
1 parent 5d07be6 commit a87dce0

26 files changed

+403
-183
lines changed

chrome/browser/ai/ai_language_model.cc

Lines changed: 49 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include "mojo/public/cpp/bindings/message.h"
3434
#include "services/on_device_model/public/cpp/capabilities.h"
3535
#include "third_party/blink/public/common/features_generated.h"
36+
#include "third_party/blink/public/mojom/ai/ai_common.mojom-shared.h"
3637
#include "third_party/blink/public/mojom/ai/ai_language_model.mojom-forward.h"
3738
#include "third_party/blink/public/mojom/ai/ai_language_model.mojom-shared.h"
3839
#include "third_party/blink/public/mojom/ai/ai_manager.mojom-shared.h"
@@ -170,9 +171,11 @@ class AILanguageModel::PromptState
170171
weak_factory_.GetWeakPtr(), std::move(session)));
171172
}
172173

173-
void OnError(blink::mojom::ModelStreamingResponseStatus error) {
174+
void OnError(blink::mojom::ModelStreamingResponseStatus error,
175+
blink::mojom::QuotaErrorInfoPtr quota_error_info = nullptr) {
174176
if (responder_) {
175-
responder_->OnError(error);
177+
AIUtils::SendStreamingStatus(responder_, error,
178+
std::move(quota_error_info));
176179
}
177180
session_.reset();
178181
responder_.reset();
@@ -502,10 +505,11 @@ void AILanguageModel::Initialize(
502505
} else {
503506
auto input = ConvertToInput(initial_prompts, session_params_->capabilities);
504507
if (!input) {
505-
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>(
506-
std::move(create_client))
507-
->OnError(
508-
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
508+
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>
509+
client_remote(std::move(create_client));
510+
AIUtils::SendClientRemoteError(
511+
client_remote,
512+
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
509513
return;
510514
}
511515
// This does not need to be queued because the AILanguageModel receiver has
@@ -625,19 +629,23 @@ void AILanguageModel::InitializeGetInputSizeComplete(
625629
create_client,
626630
std::optional<uint32_t> token_count) {
627631
if (!initial_session_ || !token_count) {
628-
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>(
629-
std::move(create_client))
630-
->OnError(
631-
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
632+
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>
633+
client_remote(std::move(create_client));
634+
AIUtils::SendClientRemoteError(
635+
client_remote,
636+
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
632637
return;
633638
}
634639

635640
uint32_t max_tokens = context_->max_tokens();
636641
if (*token_count > max_tokens) {
637-
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>(
638-
std::move(create_client))
639-
->OnError(
640-
blink::mojom::AIManagerCreateClientError::kInitialInputTooLarge);
642+
auto quota = context_->max_tokens() - context_->current_tokens();
643+
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>
644+
client_remote(std::move(create_client));
645+
AIUtils::SendClientRemoteError(
646+
client_remote,
647+
blink::mojom::AIManagerCreateClientError::kInitialInputTooLarge,
648+
blink::mojom::QuotaErrorInfo::New(token_count.value(), quota));
641649
return;
642650
}
643651

@@ -677,7 +685,8 @@ void AILanguageModel::InitializeSafetyChecksComplete(
677685
// failure.
678686
if (safety_result.failed_to_run || safety_result.is_unsafe ||
679687
safety_result.is_unsupported_language) {
680-
client->OnError(
688+
AIUtils::SendClientRemoteError(
689+
client,
681690
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
682691
return;
683692
}
@@ -698,7 +707,8 @@ void AILanguageModel::ForkInternal(
698707
mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient> remote(
699708
std::move(client));
700709
if (!initial_session_ || !model_client_) {
701-
remote->OnError(
710+
AIUtils::SendClientRemoteError(
711+
remote,
702712
blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
703713
return;
704714
}
@@ -723,19 +733,21 @@ void AILanguageModel::PromptInternal(
723733
pending_responder,
724734
base::OnceClosure on_complete) {
725735
if (!initial_session_) {
726-
mojo::Remote<blink::mojom::ModelStreamingResponder>(
727-
std::move(pending_responder))
728-
->OnError(
729-
blink::mojom::ModelStreamingResponseStatus::kErrorSessionDestroyed);
736+
mojo::Remote<blink::mojom::ModelStreamingResponder> responder(
737+
std::move(pending_responder));
738+
AIUtils::SendStreamingStatus(
739+
responder,
740+
blink::mojom::ModelStreamingResponseStatus::kErrorSessionDestroyed);
730741
return;
731742
}
732743

733744
auto input = ConvertToInputForExecute(prompts, session_params_->capabilities);
734745
if (!input) {
735-
mojo::Remote<blink::mojom::ModelStreamingResponder>(
736-
std::move(pending_responder))
737-
->OnError(
738-
blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
746+
mojo::Remote<blink::mojom::ModelStreamingResponder> responder(
747+
std::move(pending_responder));
748+
AIUtils::SendStreamingStatus(
749+
responder,
750+
blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
739751
return;
740752
}
741753
prompt_state_ = std::make_unique<PromptState>(
@@ -765,8 +777,10 @@ void AILanguageModel::PromptGetInputSizeComplete(
765777

766778
auto space_reserved = context_->ReserveSpace(*token_count);
767779
if (space_reserved == Context::SpaceReservationResult::kInsufficientSpace) {
780+
auto quota = context_->max_tokens() - context_->current_tokens();
768781
prompt_state_->OnError(
769-
blink::mojom::ModelStreamingResponseStatus::kErrorInputTooLarge);
782+
blink::mojom::ModelStreamingResponseStatus::kErrorInputTooLarge,
783+
blink::mojom::QuotaErrorInfo::New(token_count.value(), quota));
770784
return;
771785
}
772786

@@ -836,19 +850,21 @@ void AILanguageModel::AppendInternal(
836850
pending_responder,
837851
base::OnceClosure on_complete) {
838852
if (!initial_session_) {
839-
mojo::Remote<blink::mojom::ModelStreamingResponder>(
840-
std::move(pending_responder))
841-
->OnError(
842-
blink::mojom::ModelStreamingResponseStatus::kErrorSessionDestroyed);
853+
mojo::Remote<blink::mojom::ModelStreamingResponder> responder(
854+
std::move(pending_responder));
855+
AIUtils::SendStreamingStatus(
856+
responder,
857+
blink::mojom::ModelStreamingResponseStatus::kErrorSessionDestroyed);
843858
return;
844859
}
845860

846861
auto input = ConvertToInput(prompts, session_params_->capabilities);
847862
if (!input) {
848-
mojo::Remote<blink::mojom::ModelStreamingResponder>(
849-
std::move(pending_responder))
850-
->OnError(
851-
blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
863+
mojo::Remote<blink::mojom::ModelStreamingResponder> responder(
864+
std::move(pending_responder));
865+
AIUtils::SendStreamingStatus(
866+
responder,
867+
blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
852868
return;
853869
}
854870
prompt_state_ = std::make_unique<PromptState>(

chrome/browser/ai/ai_language_model_unittest.cc

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -182,13 +182,19 @@ class TestStreamingResponder : public blink::mojom::ModelStreamingResponder {
182182
return *error_status_;
183183
}
184184

185+
blink::mojom::QuotaErrorInfo quota_error_info() const {
186+
return *quota_error_info_;
187+
}
188+
185189
const std::vector<std::string> responses() const { return responses_; }
186190
uint64_t current_tokens() const { return current_tokens_; }
187191

188192
private:
189193
// blink::mojom::ModelStreamingResponder:
190-
void OnError(blink::mojom::ModelStreamingResponseStatus status) override {
194+
void OnError(blink::mojom::ModelStreamingResponseStatus status,
195+
blink::mojom::QuotaErrorInfoPtr quota_error_info) override {
191196
error_status_ = status;
197+
quota_error_info_ = std::move(quota_error_info);
192198
run_loop_.Quit();
193199
}
194200

@@ -205,6 +211,7 @@ class TestStreamingResponder : public blink::mojom::ModelStreamingResponder {
205211
void OnQuotaOverflow() override { quota_overflow_run_loop_.Quit(); }
206212

207213
std::optional<blink::mojom::ModelStreamingResponseStatus> error_status_;
214+
blink::mojom::QuotaErrorInfoPtr quota_error_info_;
208215
std::vector<std::string> responses_;
209216
uint64_t current_tokens_ = 0;
210217
base::RunLoop run_loop_;
@@ -565,20 +572,28 @@ TEST_F(AILanguageModelTest, InitialPromptsInstanceInfo) {
565572
}
566573

567574
TEST_F(AILanguageModelTest, InitialPromptsTooLarge) {
568-
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
575+
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> error_future;
576+
base::test::TestFuture<blink::mojom::QuotaErrorInfoPtr>
577+
quota_error_info_future;
569578
AITestUtils::MockCreateLanguageModelClient language_model_client;
570-
EXPECT_CALL(language_model_client, OnError(_)).WillOnce([&](auto error) {
571-
future.SetValue(error);
572-
});
579+
EXPECT_CALL(language_model_client, OnError(_, _))
580+
.WillOnce([&](auto error, auto quota_error_info) {
581+
error_future.SetValue(error);
582+
quota_error_info_future.SetValue(std::move(quota_error_info));
583+
});
573584

574585
auto options = blink::mojom::AILanguageModelCreateOptions::New();
575586
options->initial_prompts.push_back(
576587
MakePrompt(Role::kSystem, std::string(kTestMaxTokens + 1, 'a')));
577588

578589
GetAIManagerRemote()->CreateLanguageModel(
579590
language_model_client.BindNewPipeAndPassRemote(), std::move(options));
580-
EXPECT_EQ(future.Take(),
591+
EXPECT_EQ(error_future.Take(),
581592
blink::mojom::AIManagerCreateClientError::kInitialInputTooLarge);
593+
auto quota_error_info = quota_error_info_future.Take();
594+
ASSERT_TRUE(quota_error_info);
595+
ASSERT_GT(quota_error_info->requested, kTestMaxTokens);
596+
ASSERT_EQ(quota_error_info->quota, kTestMaxTokens);
582597
}
583598

584599
TEST_F(AILanguageModelTest, InputTooLarge) {
@@ -590,6 +605,8 @@ TEST_F(AILanguageModelTest, InputTooLarge) {
590605
EXPECT_FALSE(responder.WaitForCompletion());
591606
EXPECT_EQ(responder.error_status(),
592607
blink::mojom::ModelStreamingResponseStatus::kErrorInputTooLarge);
608+
ASSERT_GT(responder.quota_error_info().requested, kTestMaxTokens);
609+
ASSERT_EQ(responder.quota_error_info().quota, kTestMaxTokens);
593610
}
594611

595612
TEST_F(AILanguageModelTest, QuotaOverflowOnPromptInput) {
@@ -697,9 +714,9 @@ TEST_F(AILanguageModelTest, DestroyWithActivePrompt) {
697714
TEST_F(AILanguageModelTest, UnsupportedLanguage) {
698715
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
699716
AITestUtils::MockCreateLanguageModelClient language_model_client;
700-
EXPECT_CALL(language_model_client, OnError(_)).WillOnce([&](auto error) {
701-
future.SetValue(error);
702-
});
717+
EXPECT_CALL(language_model_client, OnError(_, _))
718+
.WillOnce(
719+
[&](auto error, auto quota_error_info) { future.SetValue(error); });
703720

704721
auto expected_input = blink::mojom::AILanguageModelExpected::New();
705722
expected_input->languages.emplace();
@@ -720,9 +737,9 @@ TEST_F(AILanguageModelTest, UnsupportedInputCapability) {
720737

721738
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
722739
AITestUtils::MockCreateLanguageModelClient language_model_client;
723-
EXPECT_CALL(language_model_client, OnError(_)).WillOnce([&](auto error) {
724-
future.SetValue(error);
725-
});
740+
EXPECT_CALL(language_model_client, OnError(_, _))
741+
.WillOnce(
742+
[&](auto error, auto quota_error_info) { future.SetValue(error); });
726743

727744
auto expected_input = blink::mojom::AILanguageModelExpected::New();
728745
expected_input->type = blink::mojom::AILanguageModelPromptType::kImage;
@@ -742,9 +759,9 @@ TEST_F(AILanguageModelTest, UnsupportedOutputCapability) {
742759

743760
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
744761
AITestUtils::MockCreateLanguageModelClient language_model_client;
745-
EXPECT_CALL(language_model_client, OnError(_)).WillOnce([&](auto error) {
746-
future.SetValue(error);
747-
});
762+
EXPECT_CALL(language_model_client, OnError(_, _))
763+
.WillOnce(
764+
[&](auto error, auto quota_error_info) { future.SetValue(error); });
748765

749766
auto expected_output = blink::mojom::AILanguageModelExpected::New();
750767
expected_output->type = blink::mojom::AILanguageModelPromptType::kImage;
@@ -904,9 +921,9 @@ TEST_F(AILanguageModelTest, TextSafetyInitialPrompts) {
904921

905922
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
906923
AITestUtils::MockCreateLanguageModelClient language_model_client;
907-
EXPECT_CALL(language_model_client, OnError(_)).WillOnce([&](auto error) {
908-
future.SetValue(error);
909-
});
924+
EXPECT_CALL(language_model_client, OnError(_, _))
925+
.WillOnce(
926+
[&](auto error, auto quota_error_info) { future.SetValue(error); });
910927

911928
auto options = blink::mojom::AILanguageModelCreateOptions::New();
912929
options->initial_prompts.push_back(MakePrompt(Role::kSystem, "unsafe"));

0 commit comments

Comments
 (0)