@@ -182,13 +182,19 @@ class TestStreamingResponder : public blink::mojom::ModelStreamingResponder {
182
182
return *error_status_;
183
183
}
184
184
185
+ blink::mojom::QuotaErrorInfo quota_error_info () const {
186
+ return *quota_error_info_;
187
+ }
188
+
185
189
const std::vector<std::string> responses () const { return responses_; }
186
190
uint64_t current_tokens () const { return current_tokens_; }
187
191
188
192
private:
189
193
// 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 {
191
196
error_status_ = status;
197
+ quota_error_info_ = std::move (quota_error_info);
192
198
run_loop_.Quit ();
193
199
}
194
200
@@ -205,6 +211,7 @@ class TestStreamingResponder : public blink::mojom::ModelStreamingResponder {
205
211
void OnQuotaOverflow () override { quota_overflow_run_loop_.Quit (); }
206
212
207
213
std::optional<blink::mojom::ModelStreamingResponseStatus> error_status_;
214
+ blink::mojom::QuotaErrorInfoPtr quota_error_info_;
208
215
std::vector<std::string> responses_;
209
216
uint64_t current_tokens_ = 0 ;
210
217
base::RunLoop run_loop_;
@@ -565,20 +572,28 @@ TEST_F(AILanguageModelTest, InitialPromptsInstanceInfo) {
565
572
}
566
573
567
574
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;
569
578
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
+ });
573
584
574
585
auto options = blink::mojom::AILanguageModelCreateOptions::New ();
575
586
options->initial_prompts .push_back (
576
587
MakePrompt (Role::kSystem , std::string (kTestMaxTokens + 1 , ' a' )));
577
588
578
589
GetAIManagerRemote ()->CreateLanguageModel (
579
590
language_model_client.BindNewPipeAndPassRemote (), std::move (options));
580
- EXPECT_EQ (future .Take (),
591
+ EXPECT_EQ (error_future .Take (),
581
592
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 );
582
597
}
583
598
584
599
TEST_F (AILanguageModelTest, InputTooLarge) {
@@ -590,6 +605,8 @@ TEST_F(AILanguageModelTest, InputTooLarge) {
590
605
EXPECT_FALSE (responder.WaitForCompletion ());
591
606
EXPECT_EQ (responder.error_status (),
592
607
blink::mojom::ModelStreamingResponseStatus::kErrorInputTooLarge );
608
+ ASSERT_GT (responder.quota_error_info ().requested , kTestMaxTokens );
609
+ ASSERT_EQ (responder.quota_error_info ().quota , kTestMaxTokens );
593
610
}
594
611
595
612
TEST_F (AILanguageModelTest, QuotaOverflowOnPromptInput) {
@@ -697,9 +714,9 @@ TEST_F(AILanguageModelTest, DestroyWithActivePrompt) {
697
714
TEST_F (AILanguageModelTest, UnsupportedLanguage) {
698
715
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
699
716
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); });
703
720
704
721
auto expected_input = blink::mojom::AILanguageModelExpected::New ();
705
722
expected_input->languages .emplace ();
@@ -720,9 +737,9 @@ TEST_F(AILanguageModelTest, UnsupportedInputCapability) {
720
737
721
738
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
722
739
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); });
726
743
727
744
auto expected_input = blink::mojom::AILanguageModelExpected::New ();
728
745
expected_input->type = blink::mojom::AILanguageModelPromptType::kImage ;
@@ -742,9 +759,9 @@ TEST_F(AILanguageModelTest, UnsupportedOutputCapability) {
742
759
743
760
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
744
761
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); });
748
765
749
766
auto expected_output = blink::mojom::AILanguageModelExpected::New ();
750
767
expected_output->type = blink::mojom::AILanguageModelPromptType::kImage ;
@@ -904,9 +921,9 @@ TEST_F(AILanguageModelTest, TextSafetyInitialPrompts) {
904
921
905
922
base::test::TestFuture<blink::mojom::AIManagerCreateClientError> future;
906
923
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); });
910
927
911
928
auto options = blink::mojom::AILanguageModelCreateOptions::New ();
912
929
options->initial_prompts .push_back (MakePrompt (Role::kSystem , " unsafe" ));
0 commit comments