From 205b57ec1eb7b9ab3bdb34a5666f5c229b5443e8 Mon Sep 17 00:00:00 2001
From: Mike Wasserman <[email protected]>
Date: Wed, 20 May 2026 13:33:40 -0700
Subject: [PATCH] Prompt API: Gate tool use requests by feature flag in browser

Add browser-side enforcement of renderer-side feature checks.
Add unit tests.

Fixed: 513271007
Change-Id: I162511571020fdb8abb55f9354d6d9fd6c28d8d2
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7858163
Reviewed-by: Isaac Ahouma <[email protected]>
Auto-Submit: Mike Wasserman <[email protected]>
Commit-Queue: Mike Wasserman <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1633802}
---

diff --git a/chrome/browser/ai/ai_language_model.cc b/chrome/browser/ai/ai_language_model.cc
index a35583d2..9a715f6b 100644
--- a/chrome/browser/ai/ai_language_model.cc
+++ b/chrome/browser/ai/ai_language_model.cc
@@ -159,6 +159,9 @@
           // TODO(crbug.com/422803232): Support on_device_model tool use.
           return nullptr;
         case blink::mojom::AILanguageModelPromptContent::Tag::kToolResponse: {
+          if (!capabilities.Has(on_device_model::CapabilityFlags::kToolUse)) {
+            return nullptr;
+          }
           // TODO(crbug.com/422803232): Support image/audio result types.
           auto response = GetToolResponseFromDict(content->get_tool_response());
           if (!response) {
diff --git a/chrome/browser/ai/ai_language_model_unittest.cc b/chrome/browser/ai/ai_language_model_unittest.cc
index db0fdf45..3a9f58e 100644
--- a/chrome/browser/ai/ai_language_model_unittest.cc
+++ b/chrome/browser/ai/ai_language_model_unittest.cc
@@ -1827,6 +1827,12 @@
 
 // Test class for `Tool Use` functionality.
 class AILanguageModelOpenLoopToolTest : public AILanguageModelTest {
+ public:
+  AILanguageModelOpenLoopToolTest() {
+    scoped_feature_list_.InitAndEnableFeature(
+        blink::features::kAIPromptAPIToolUse);
+  }
+
  protected:
   // Override CreateConfig to use higher max tokens for `Tool Use` testing.
   optimization_guide::proto::OnDeviceModelExecutionFeatureConfig CreateConfig()
@@ -1903,6 +1909,8 @@
   void DisableToolCallSimulation() {
     fake_broker_->settings().simulated_tool_calls.clear();
   }
+
+  base::test::ScopedFeatureList scoped_feature_list_;
 };
 
 // Test that tools are embedded in the session's initial context.
@@ -2274,6 +2282,58 @@
   EXPECT_THAT(final_response, testing::HasSubstr("\"condition\":\"Cloudy\""));
 }
 
+// Verify that CanCreate and Create correctly handle the disabled tool use flag.
+TEST_F(AILanguageModelOpenLoopToolTest, RejectCreateWithFlagDisabled) {
+  base::test::ScopedFeatureList scoped_feature_list;
+  scoped_feature_list.InitAndDisableFeature(
+      blink::features::kAIPromptAPIToolUse);
+
+  std::vector<blink::mojom::AILanguageModelToolDeclarationPtr> tools;
+  tools.push_back(CreateWeatherTool());
+  auto options = blink::mojom::AILanguageModelCreateOptions::New();
+  options->tools = std::move(tools);
+
+  base::test::TestFuture<blink::mojom::ModelAvailabilityCheckResult> future;
+  GetAIManagerInterface()->CanCreateLanguageModel(options.Clone(),
+                                                  future.GetCallback());
+  EXPECT_EQ(future.Get(), blink::mojom::ModelAvailabilityCheckResult::
+                              kUnavailableModelAdaptationNotAvailable);
+
+  TestCreateLanguageModelClient client;
+  GetAIManagerRemote()->CreateLanguageModel(client.BindNewPipeAndPassRemote(),
+                                            std::move(options));
+  auto result = client.result().Take();
+  ASSERT_FALSE(result.has_value());
+  EXPECT_EQ(result.error().error,
+            blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
+}
+
+// Ensure Prompt rejects tool responses on sessions without the tool capability.
+TEST_F(AILanguageModelOpenLoopToolTest, RejectToolResponseWithoutCapability) {
+  auto session = CreateSession();
+  base::DictValue tool_response_dict;
+  tool_response_dict.Set("callID", "c0");
+  tool_response_dict.Set("name", "no_such_tool");
+  base::DictValue result;
+  result.Set("arbitrary_field", "arbitrary_value");
+  tool_response_dict.Set("result", std::move(result));
+
+  std::vector<blink::mojom::AILanguageModelPromptPtr> prompts;
+  auto prompt = blink::mojom::AILanguageModelPrompt::New();
+  prompt->role = Role::kUser;
+  prompt->content.push_back(
+      blink::mojom::AILanguageModelPromptContent::NewToolResponse(
+          std::move(tool_response_dict)));
+  prompts.push_back(std::move(prompt));
+
+  AITestUtils::TestStreamingResponder responder;
+  session->Prompt(std::move(prompts), /*constraint=*/nullptr,
+                  responder.BindRemote());
+  ASSERT_FALSE(responder.WaitForCompletion());
+  EXPECT_EQ(responder.error_status(),
+            blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
+}
+
 TEST_F(AILanguageModelTest, CanCreatePermissionsPolicyDisabled) {
   DisablePolicy(network::mojom::PermissionsPolicyFeature::kLanguageModel);
   mojo::test::BadMessageObserver observer;
diff --git a/chrome/browser/ai/ai_manager.cc b/chrome/browser/ai/ai_manager.cc
index 7539131..abf25e5 100644
--- a/chrome/browser/ai/ai_manager.cc
+++ b/chrome/browser/ai/ai_manager.cc
@@ -704,9 +704,14 @@
                                   kUnavailableModelAdaptationNotAvailable);
       return;
     }
-    // Note: Tool use capabilities are gated by RuntimeEnabledFeatures in Blink.
     // Tool use capability is signaled by the presence of tool declarations.
     if (options->tools.has_value() && !options->tools->empty()) {
+      // Check if tool declarations are used without the feature flag.
+      if (!base::FeatureList::IsEnabled(blink::features::kAIPromptAPIToolUse)) {
+        std::move(callback).Run(blink::mojom::ModelAvailabilityCheckResult::
+                                    kUnavailableModelAdaptationNotAvailable);
+        return;
+      }
       input_capabilities.Put(on_device_model::CapabilityFlags::kToolUse);
     }
   }
@@ -865,11 +870,13 @@
     return;
   }
   if (!params->capabilities.empty()) {
-    // Check if multimodal input (image/audio) is used without the feature flag
-    // or if the model doesn't support the requested capabilities.
+    // Check if image/audio/tool input types are used without feature flags, or
+    // the model doesn't support the requested capabilities.
     if ((HasMultimodalInputCapabilities(params->capabilities) &&
          !base::FeatureList::IsEnabled(
              blink::features::kAIPromptAPIMultimodalInput)) ||
+        (params->capabilities.Has(on_device_model::CapabilityFlags::kToolUse) &&
+         !base::FeatureList::IsEnabled(blink::features::kAIPromptAPIToolUse)) ||
         !model_client->capabilities().HasAll(params->capabilities)) {
       mojo::Remote<blink::mojom::AIManagerCreateLanguageModelClient>
           client_remote(std::move(client));

diff --git a/chrome/browser/ai/ai_language_model_unittest.cc b/chrome/browser/ai/ai_language_model_unittest.cc
index db0fdf45..3a9f58e 100644
--- a/chrome/browser/ai/ai_language_model_unittest.cc
+++ b/chrome/browser/ai/ai_language_model_unittest.cc
@@ -1827,6 +1827,12 @@
 
 // Test class for `Tool Use` functionality.
 class AILanguageModelOpenLoopToolTest : public AILanguageModelTest {
+ public:
+  AILanguageModelOpenLoopToolTest() {
+    scoped_feature_list_.InitAndEnableFeature(
+        blink::features::kAIPromptAPIToolUse);
+  }
+
  protected:
   // Override CreateConfig to use higher max tokens for `Tool Use` testing.
   optimization_guide::proto::OnDeviceModelExecutionFeatureConfig CreateConfig()
@@ -1903,6 +1909,8 @@
   void DisableToolCallSimulation() {
     fake_broker_->settings().simulated_tool_calls.clear();
   }
+
+  base::test::ScopedFeatureList scoped_feature_list_;
 };
 
 // Test that tools are embedded in the session's initial context.
@@ -2274,6 +2282,58 @@
   EXPECT_THAT(final_response, testing::HasSubstr("\"condition\":\"Cloudy\""));
 }
 
+// Verify that CanCreate and Create correctly handle the disabled tool use flag.
+TEST_F(AILanguageModelOpenLoopToolTest, RejectCreateWithFlagDisabled) {
+  base::test::ScopedFeatureList scoped_feature_list;
+  scoped_feature_list.InitAndDisableFeature(
+      blink::features::kAIPromptAPIToolUse);
+
+  std::vector<blink::mojom::AILanguageModelToolDeclarationPtr> tools;
+  tools.push_back(CreateWeatherTool());
+  auto options = blink::mojom::AILanguageModelCreateOptions::New();
+  options->tools = std::move(tools);
+
+  base::test::TestFuture<blink::mojom::ModelAvailabilityCheckResult> future;
+  GetAIManagerInterface()->CanCreateLanguageModel(options.Clone(),
+                                                  future.GetCallback());
+  EXPECT_EQ(future.Get(), blink::mojom::ModelAvailabilityCheckResult::
+                              kUnavailableModelAdaptationNotAvailable);
+
+  TestCreateLanguageModelClient client;
+  GetAIManagerRemote()->CreateLanguageModel(client.BindNewPipeAndPassRemote(),
+                                            std::move(options));
+  auto result = client.result().Take();
+  ASSERT_FALSE(result.has_value());
+  EXPECT_EQ(result.error().error,
+            blink::mojom::AIManagerCreateClientError::kUnableToCreateSession);
+}
+
+// Ensure Prompt rejects tool responses on sessions without the tool capability.
+TEST_F(AILanguageModelOpenLoopToolTest, RejectToolResponseWithoutCapability) {
+  auto session = CreateSession();
+  base::DictValue tool_response_dict;
+  tool_response_dict.Set("callID", "c0");
+  tool_response_dict.Set("name", "no_such_tool");
+  base::DictValue result;
+  result.Set("arbitrary_field", "arbitrary_value");
+  tool_response_dict.Set("result", std::move(result));
+
+  std::vector<blink::mojom::AILanguageModelPromptPtr> prompts;
+  auto prompt = blink::mojom::AILanguageModelPrompt::New();
+  prompt->role = Role::kUser;
+  prompt->content.push_back(
+      blink::mojom::AILanguageModelPromptContent::NewToolResponse(
+          std::move(tool_response_dict)));
+  prompts.push_back(std::move(prompt));
+
+  AITestUtils::TestStreamingResponder responder;
+  session->Prompt(std::move(prompts), /*constraint=*/nullptr,
+                  responder.BindRemote());
+  ASSERT_FALSE(responder.WaitForCompletion());
+  EXPECT_EQ(responder.error_status(),
+            blink::mojom::ModelStreamingResponseStatus::kErrorInvalidRequest);
+}
+
 TEST_F(AILanguageModelTest, CanCreatePermissionsPolicyDisabled) {
   DisablePolicy(network::mojom::PermissionsPolicyFeature::kLanguageModel);
   mojo::test::BadMessageObserver observer;