chrome/browser/extensions/ai_language_model_browsertest.cc

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <tuple>
#include <vector>

#include "base/command_line.h"
#include "base/strings/to_string.h"
#include "base/test/scoped_feature_list.h"
#include "base/version_info/channel.h"
#include "chrome/browser/extensions/extension_browsertest.h"
#include "components/embedder_support/switches.h"
#include "content/public/common/content_switches.h"
#include "content/public/test/browser_test.h"
#include "extensions/test/result_catcher.h"
#include "extensions/test/test_extension_dir.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/blink/public/common/features_generated.h"

// TODO(crbug.com/350642260): the prompt API for extension OT is not affecting
// ChromeOS. We have skipped the logic for ChromeOS so the test will be skipped
// as well.
namespace extensions {

namespace {

// This is the public key of tools/origin_trials/eftest.key, used to validate
// origin trial tokens generated by tools/origin_trials/generate_token.py.
constexpr char kOriginTrialPublicKeyForTesting[] =
    "dRCs+TocuKkocNKa0AtZ4awrt9XKH2SQCI6o4FY6BNA=";

// The extension origin trial token (expired on 2032-11-26) was generated by
// ```
// tools/origin_trials/generate_token.py
// chrome-extension://jnapclmfkaejhjkddbmiafekigmcbmma AIPromptAPIForExtension
// --expire-days 3000
// ```
constexpr char kLanguageModelOriginTrialTokensField[] =
    "\"trial_tokens\":[\"A5nDxhrF7Qe4GiLouR1mgL5XKSk4wXA0B/RV2VyQcZj2IkLALdG/"
    "FHrucKbG1TKD8QidNfqBdC07wP8KJaF6EQYAAAB9eyJvcmlnaW4iOiAiY2hyb21lLWV4dGVuc2"
    "lvbjovL2puYXBjbG1ma2Flamhqa2RkYm1pYWZla2lnbWNibW1hIiwgImZlYXR1cmUiOiAiQUlQ"
    "cm9tcHRBUElGb3JFeHRlbnNpb24iLCAiZXhwaXJ5IjogMTk4NTA2MjMwM30=\"],";

// The `key` field stores the public key for the extension with id
// "jnapclmfkaejhjkddbmiafekigmcbmma".
static constexpr char kManifestTemplate[] =
    R"JS(
    {
      "name": "AI language model test",
      "version": "0.1",
      "manifest_version": 3,
      %s
      "key": "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA3H6Jc0On6l0H3DJ6bx4aOW3+srCfjSdr+3ukwIEZrL6jDy500XweIwOp9PItpM9sijwu8v1rdyoBPubm/ottp/oz42aKp+2xIxcMTa6/cA2BL2kOWxwv+WP9d01IOFbFpWmQBDQNpp2UmH67OFbie6zHhyrSJKL2o9d05iX0a9Xwv9W48JKYpldo+/2JTP/5en0jxgiN+qkOCZuLag2cS/6Az0LArqsf5D+ReJemIBCNJhVxu3P0naxfEG6B6XczzuuptrX3H2vDr1LxZasLh9bzV88+8BxarjETACebOfqy366QxXluwAjnu/NHPv53edXlXvXrZ0C69RvvlMh1qQIDAQAB",
      "description": "Extension for testing the AI language model API.",
      "background": {
        "service_worker": "sw.js"
      }
    }
  )JS";

static constexpr char kServiceWorkerScript[] =
    R"JS(
      chrome.test.runTests([
        function verifyLanguageModel() {
          const expectLanguageModel = %s;
          chrome.test.assertEq(expectLanguageModel, !!self.LanguageModel);
          chrome.test.succeed();
        },
      ]);
    )JS";

// The boolean tuple describing:
// 1. if the chrome://flag `kAIPromptAPI` is explicitly enabled;
// 2. if the kAIPromptAPI kill switch is triggered;
// 3. if the chrome://flag `kAIPromptAPIForExtension` is explicitly enabled;
// 4. if the kAIPromptAPIForExtension kill switch is triggered;
// 5. if the extension has an origin trial token.
using Variant = std::tuple<bool, bool, bool, bool, bool>;
bool IsAPIFlagEnabled(Variant v) {
  return std::get<0>(v);
}
bool IsAPIKillSwitchTriggered(Variant v) {
  return std::get<1>(v);
}
bool IsExtensionFlagEnabled(Variant v) {
  return std::get<2>(v);
}
bool IsExtensionKillSwitchTriggered(Variant v) {
  return std::get<3>(v);
}
bool IsExtensionInOriginTrial(Variant v) {
  return std::get<4>(v);
}

// Describes the test variants in a meaningful way in the parameterized tests.
std::string DescribeTestVariant(const testing::TestParamInfo<Variant> info) {
  std::string api_flag_enabled =
      IsAPIFlagEnabled(info.param) ? "WithAPIFlag" : "NoAPIFlag";
  std::string api_kill_switch = IsAPIKillSwitchTriggered(info.param)
                                    ? "WithAPIKillswitch"
                                    : "NoAPIKillswitch";
  std::string extension_flag_enabled = IsExtensionFlagEnabled(info.param)
                                           ? "WithExtensionFlag"
                                           : "NoExtensionFlag";
  std::string extension_kill_switch = IsExtensionKillSwitchTriggered(info.param)
                                          ? "WithExtensionKillswitch"
                                          : "NoExtensionKillswitch";
  std::string origin_trial =
      IsExtensionInOriginTrial(info.param) ? "WithOTToken" : "NoOTToken";
  return base::JoinString(
      {api_flag_enabled, api_kill_switch, extension_flag_enabled,
       extension_kill_switch, origin_trial},
      "_");
}

}  // namespace

class ExtensionAILanguageModelBrowserTest
    : public ExtensionBrowserTest,
      public testing::WithParamInterface<Variant> {
 public:
  void SetUpCommandLine(base::CommandLine* command_line) override {
    ExtensionBrowserTest::SetUpCommandLine(command_line);
    if (IsAPIFlagEnabled(GetParam())) {
      command_line->AppendSwitchASCII(switches::kEnableBlinkFeatures,
                                      "AIPromptAPI");
    }
    if (IsExtensionFlagEnabled(GetParam())) {
      command_line->AppendSwitchASCII(switches::kEnableBlinkFeatures,
                                      "AIPromptAPIForExtension");
    }

    // Also specify the test public key to make the test token effective.
    command_line->AppendSwitchASCII(embedder_support::kOriginTrialPublicKey,
                                    kOriginTrialPublicKeyForTesting);

    base::flat_map<base::test::FeatureRef, bool> feature_states;
    if (IsAPIKillSwitchTriggered(GetParam())) {
      feature_states[blink::features::kAIPromptAPI] = false;
    }
    if (IsExtensionKillSwitchTriggered(GetParam())) {
      feature_states[blink::features::kAIPromptAPIForExtension] = false;
    }
    feature_list_.InitWithFeatureStates(feature_states);
  }

 protected:
  std::string GetManifest() {
    return base::StringPrintf(kManifestTemplate,
                              IsExtensionInOriginTrial(GetParam())
                                  ? kLanguageModelOriginTrialTokensField
                                  : "");
  }

 private:
  base::test::ScopedFeatureList feature_list_;
};

INSTANTIATE_TEST_SUITE_P(
    /* no prefix */,
    ExtensionAILanguageModelBrowserTest,
    testing::Combine(testing::Bool(),
                     testing::Bool(),
                     testing::Bool(),
                     testing::Bool(),
                     testing::Bool()),
    &DescribeTestVariant);

#if BUILDFLAG(IS_CHROMEOS)
#define MAYBE_WorkerAccess DISABLED_WorkerAccess
#else
#define MAYBE_WorkerAccess WorkerAccess
#endif  // BUILDFLAG(IS_CHROMEOS)
IN_PROC_BROWSER_TEST_P(ExtensionAILanguageModelBrowserTest,
                       MAYBE_WorkerAccess) {
  TestExtensionDir test_dir;
  test_dir.WriteManifest(GetManifest());
  const bool is_api_exposed = IsAPIFlagEnabled(GetParam()) ||
                              IsExtensionFlagEnabled(GetParam()) ||
                              (!IsAPIKillSwitchTriggered(GetParam()) &&
                               !IsExtensionKillSwitchTriggered(GetParam()) &&
                               IsExtensionInOriginTrial(GetParam()));
  test_dir.WriteFile(
      FILE_PATH_LITERAL("sw.js"),
      base::StringPrintf(kServiceWorkerScript, base::ToString(is_api_exposed)));
  ResultCatcher result_catcher;
  const Extension* extension = LoadExtension(test_dir.UnpackedPath());
  ASSERT_TRUE(extension);
  EXPECT_TRUE(result_catcher.GetNextResult()) << result_catcher.message();
}

}  // namespace extensions