[AI Chat] Add semantic history search tool for Leo (#36385)

* [AI Chat] Add HistorySearchTool exposing local history embeddings

Adds a self-executing Tool subclass that wraps
HistoryEmbeddingsService::Search() so the Leo assistant can semantically
search the user's local browsing history. Returns matches as a JSON
TextContentBlock containing title, URL, snippet, and last visit time.

The tool is not yet wired into any provider; that follows in a separate
commit so this change can be reviewed in isolation.

* [AI Chat] Add unit tests for HistorySearchTool

Covers the input validation and JSON serialization paths exercised by
HistorySearchTool. The full Search() flow is intentionally not faked
here -- HistoryEmbeddingsService has a heavyweight constructor and
faking it pulls in too many dependencies for a tool unit test; the
JSON formatter is exposed via internal:: so it can be tested directly.

* [AI Chat] Wire HistorySearchTool into BrowserToolProvider

The tool is created when history embeddings is enabled for the active
profile (via IsHistoryEmbeddingsEnabledForProfile, which checks the
flag and per-profile eligibility) and surfaced through GetTools().

* [AI Chat] Permit anchors for URLs returned by client-side tools

For Leo models, MarkdownRenderer strips any anchor whose href isn't in the
entry's allowedLinks list -- which today is populated only from the
sourcesEvent that web search emits. Client-side tools (e.g. semantic
history search) had no way to surface trusted URLs, so any [text](url) the
assistant wrote referencing tool results was demoted to a <span>.

Add a sidechannel via the existing ToolArtifact: a new artifact type
"trusted_links" carrying a JSON array of HTTPS URL strings. Conversation
entries flatten these into allowedLinks for the whole assistant group, so
URLs returned by a tool call (which lives in a separate assistant entry
from the follow-up response text) permit anchors in the response.

The artifact is not rendered visually -- it's purely a trust-list. No
mojom struct change, no LLM payload pollution, no "Sources" panel.

HistorySearchTool emits the artifact populated from the URLs in its JSON
results.

* [AI Chat] Add browser test for HistorySearchTool dispatch

Verifies the end-to-end path that unit tests can't reach: the mocked
engine emits a semantic_history_search tool_use event, ConversationHandler
dispatches to the real HistorySearchTool, the tool forwards the parsed
arguments to a fake HistoryEmbeddingsSearch, and the JSON-serialized
result flows back into the conversation history.

A second case exercises malformed args: the tool surfaces an error and
the search service is never called.

Uses the existing AIChatConversationUIBrowserTestBase helper.
This commit is contained in:
Anthony Tseng
2026-06-04 04:04:09 +01:00
committed by GitHub
parent bcd925ed85
commit 3461528ff6
22 changed files with 1321 additions and 158 deletions
+20 -1
View File
@@ -35,13 +35,20 @@ static_library("ai_chat") {
"tab_tracker_service_factory.h",
"tools/code_execution_tool.cc",
"tools/code_execution_tool.h",
"tools/history_search_tool.cc",
"tools/history_search_tool.h",
]
if (enable_ai_chat_tab_management_tool) {
sources += [ "tools/tab_management_tool.h" ]
}
public_deps = [ "//brave/components/restricted_web_contents_delegate" ]
public_deps = [
"//brave/components/restricted_web_contents_delegate",
# history_search_tool.h includes history_embeddings_search.h.
"//components/history_embeddings/core",
]
deps = [
"//brave/brave_domains",
@@ -61,11 +68,14 @@ static_library("ai_chat") {
"//brave/net/base:utils",
"//chrome/browser:browser_process",
"//chrome/browser/actor",
"//chrome/browser/history_embeddings",
"//chrome/browser/profiles:profile",
"//chrome/browser/ui/tabs:tabs_public",
"//chrome/common",
"//chrome/common:channel_info",
"//components/browsing_data/core",
"//components/history/core/browser",
"//components/history_embeddings/content",
"//components/keyed_service/content",
"//components/prefs",
"//components/user_prefs",
@@ -290,6 +300,7 @@ source_set("unit_tests") {
"ai_chat_throttle_unittest.cc",
"brave_open_ai_chat_permission_context_unittest.cc",
"full_screenshotter_unittest.cc",
"tools/history_search_tool_unittest.cc",
]
deps = [
@@ -302,11 +313,16 @@ source_set("unit_tests") {
"//brave/components/constants",
"//chrome/common",
"//chrome/test:test_support",
"//components/history/core/browser",
"//components/history_embeddings/content",
"//components/history_embeddings/core",
"//components/paint_preview/common:test_utils",
"//components/paint_preview/common/mojom",
"//components/passage_embeddings/core",
"//components/services/paint_preview_compositor/public/mojom",
"//content/public/browser",
"//content/test:test_support",
"//testing/gmock",
"//testing/gtest",
"//url",
]
@@ -389,6 +405,7 @@ source_set("browser_tests") {
"code_execution_tool_browsertest.cc",
"page_content_fetcher_browsertest.cc",
"text_file_extractor_browsertest.cc",
"tools/history_search_tool_browsertest.cc",
]
if (enable_pdf) {
@@ -422,6 +439,8 @@ source_set("browser_tests") {
"//chrome/browser/ui/side_panel:side_panel_views_dependent",
"//chrome/test:test_support",
"//chrome/test:test_support_ui",
"//components/history/core/browser",
"//components/history_embeddings/core",
"//printing/buildflags",
"//services/screen_ai/buildflags",
"//services/screen_ai/public/cpp:utilities",
+15
View File
@@ -8,12 +8,15 @@
#include <memory>
#include <vector>
#include "base/check_is_test.h"
#include "base/feature_list.h"
#include "base/memory/weak_ptr.h"
#include "brave/browser/ai_chat/tools/code_execution_tool.h"
#include "brave/browser/ai_chat/tools/history_search_tool.h"
#include "brave/components/ai_chat/core/browser/tools/tool.h"
#include "brave/components/ai_chat/core/common/buildflags/buildflags.h"
#include "brave/components/ai_chat/core/common/features.h"
#include "chrome/browser/history_embeddings/history_embeddings_utils.h"
#include "chrome/browser/profiles/profile.h"
#include "content/public/browser/browser_context.h"
@@ -34,6 +37,9 @@ std::vector<base::WeakPtr<Tool>> BrowserToolProvider::GetTools() {
if (code_execution_tool_) {
tool_ptrs.push_back(code_execution_tool_->GetWeakPtr());
}
if (history_search_tool_) {
tool_ptrs.push_back(history_search_tool_->GetWeakPtr());
}
#if BUILDFLAG(ENABLE_AI_CHAT_TAB_MANAGEMENT_TOOL)
if (tab_management_tool_) {
@@ -44,11 +50,20 @@ std::vector<base::WeakPtr<Tool>> BrowserToolProvider::GetTools() {
return tool_ptrs;
}
HistorySearchTool* BrowserToolProvider::GetHistorySearchToolForTesting() {
CHECK_IS_TEST();
return history_search_tool_.get();
}
void BrowserToolProvider::CreateTools(
content::BrowserContext* browser_context) {
if (features::IsCodeExecutionToolEnabled()) {
code_execution_tool_ = std::make_unique<CodeExecutionTool>(browser_context);
}
if (history_embeddings::IsHistoryEmbeddingsEnabledForProfile(
Profile::FromBrowserContext(browser_context))) {
history_search_tool_ = std::make_unique<HistorySearchTool>(browser_context);
}
#if BUILDFLAG(ENABLE_AI_CHAT_TAB_MANAGEMENT_TOOL)
if (base::FeatureList::IsEnabled(features::kTabManagementTool)) {
tab_management_tool_ = std::make_unique<TabManagementTool>(profile_);
+4
View File
@@ -23,6 +23,7 @@ class BrowserContext;
namespace ai_chat {
class CodeExecutionTool;
class HistorySearchTool;
class TabManagementTool;
// Implementation of ToolProvider that provides browser-specific
@@ -41,11 +42,14 @@ class BrowserToolProvider : public ToolProvider {
// ToolProvider implementation
std::vector<base::WeakPtr<Tool>> GetTools() override;
HistorySearchTool* GetHistorySearchToolForTesting();
private:
void CreateTools(content::BrowserContext* browser_context);
// Browser-specific tools owned by this provider
std::unique_ptr<CodeExecutionTool> code_execution_tool_;
std::unique_ptr<HistorySearchTool> history_search_tool_;
#if BUILDFLAG(ENABLE_AI_CHAT_TAB_MANAGEMENT_TOOL)
std::unique_ptr<TabManagementTool> tab_management_tool_;
#endif
+12
View File
@@ -13,3 +13,15 @@ These tools should be owned by the ContentAgentToolProvider which can own the
shared state between these tools for a single conversation, and provide
information about the current Task and available Tabs to all the related tools
via ContentAgentTaskProvider.
## Browser Tools
Standalone, single-shot tools that expose a piece of browser data or
functionality to the assistant without participating in the actor Task
lifecycle. They typically wrap a single keyed service (history, bookmarks, etc.)
and return a structured JSON payload as their tool output. Examples:
`HistorySearchTool`, `CodeExecutionTool`.
These tools should be owned by the `BrowserToolProvider`, which creates them
per-conversation gated on whatever feature/per-profile checks the underlying
capability requires.
@@ -0,0 +1,299 @@
// Copyright (c) 2026 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.
#include "brave/browser/ai_chat/tools/history_search_tool.h"
#include <algorithm>
#include <utility>
#include "base/functional/bind.h"
#include "base/i18n/time_formatting.h"
#include "base/json/json_reader.h"
#include "base/json/json_writer.h"
#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/utf_string_conversions.h"
#include "base/time/time.h"
#include "base/values.h"
#include "brave/components/ai_chat/core/browser/tools/tool_input_properties.h"
#include "brave/components/ai_chat/core/browser/tools/tool_utils.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "chrome/browser/history_embeddings/history_embeddings_service_factory.h"
#include "chrome/browser/history_embeddings/history_embeddings_utils.h"
#include "chrome/browser/profiles/profile.h"
#include "components/history/core/browser/url_row.h"
#include "components/history_embeddings/content/history_embeddings_service.h"
#include "components/history_embeddings/core/history_embeddings_features.h"
#include "url/url_constants.h"
namespace ai_chat {
namespace {
constexpr char kPropertyQuery[] = "query";
constexpr char kPropertyCount[] = "count";
constexpr char kPropertyTimeRangeStartDaysAgo[] = "time_range_start_days_ago";
constexpr char kPropertyIncludeAllPassages[] = "include_all_passages";
constexpr char kOutputKeyQuery[] = "query";
constexpr char kOutputKeyResults[] = "results";
constexpr char kOutputKeyTitle[] = "title";
constexpr char kOutputKeyUrl[] = "url";
constexpr char kOutputKeyLastVisitTime[] = "last_visit_time";
constexpr char kOutputKeyPassages[] = "passages";
constexpr char kOutputKeySnippet[] = "snippet";
constexpr size_t kDefaultResultCount = 5;
constexpr size_t kMaxResultCount = 10;
base::DictValue ScoredUrlRowToDict(const history_embeddings::ScoredUrlRow& row,
bool include_all_passages) {
base::DictValue entry;
entry.Set(kOutputKeyTitle, base::UTF16ToUTF8(row.row.title()));
entry.Set(kOutputKeyUrl, row.row.url().spec());
// GetBestPassage() / GetBestScoreIndices() CHECK that there is at least
// one passage; some rows (e.g. URL-only history entries with no embedded
// passages) won't have any, so guard against it.
const int passages_size = row.url_data.passages.passages_size();
if (passages_size > 0) {
if (include_all_passages) {
base::ListValue passages;
// Returns indices ordered by descending score.
auto indices = row.GetBestScoreIndices(/*min_count=*/passages_size,
/*min_word_count=*/0);
for (size_t i : indices) {
passages.Append(row.url_data.passages.passages(i));
}
entry.Set(kOutputKeyPassages, std::move(passages));
} else {
entry.Set(kOutputKeySnippet, row.GetBestPassage());
}
}
if (!row.row.last_visit().is_null()) {
entry.Set(kOutputKeyLastVisitTime,
base::TimeFormatAsIso8601(row.row.last_visit()));
}
return entry;
}
} // namespace
namespace internal {
std::string BuildHistorySearchResultJson(
const std::string& query,
const history_embeddings::SearchResult& result,
bool include_all_passages) {
base::ListValue results;
for (const auto& row : result.scored_url_rows) {
results.Append(ScoredUrlRowToDict(row, include_all_passages));
}
base::DictValue root;
root.Set(kOutputKeyQuery, query);
root.Set(kOutputKeyResults, std::move(results));
std::string json;
base::JSONWriter::Write(root, &json);
return json;
}
} // namespace internal
HistorySearchTool::HistorySearchTool(content::BrowserContext* browser_context)
: profile_(Profile::FromBrowserContext(browser_context)) {}
HistorySearchTool::~HistorySearchTool() = default;
std::string_view HistorySearchTool::Name() const {
return mojom::kSemanticHistorySearchToolName;
}
std::string_view HistorySearchTool::Description() const {
return "Performs a semantic (meaning-based) search over the user's local "
"browsing history. Use when the user asks about something they "
"previously read but can't recall the URL or title, or when you "
"need the page's text content to answer their question. Returns "
"matching pages as JSON with title, URL, last visit time, and "
"(by default) the indexed text passages -- the only way to "
"access page content from history; you do not need to ask the "
"user to fetch the page. Do not use for exact URL/title lookup "
"or general web search. Runs entirely on-device.";
}
std::optional<base::DictValue> HistorySearchTool::InputProperties() const {
return CreateInputProperties(
{{kPropertyQuery,
StringProperty("Natural language description of the page the user "
"is trying to find in their browsing history.")},
{kPropertyCount,
IntegerProperty(
"Maximum number of results to return. Defaults to 5; capped "
"at 10.")},
{kPropertyTimeRangeStartDaysAgo,
IntegerProperty(
"If set, restrict the search to pages visited within the last "
"N days. Omit to search the entire history.")},
{kPropertyIncludeAllPassages,
BooleanProperty(
"Include all indexed passages per result, sorted by "
"relevance. Defaults to true. Set false to return a single "
"'snippet' (best-matching passage only) when no page "
"content is needed.")}});
}
std::optional<std::vector<std::string>> HistorySearchTool::RequiredProperties()
const {
return std::vector<std::string>{kPropertyQuery};
}
std::variant<bool, mojom::PermissionChallengePtr>
HistorySearchTool::RequiresUserInteractionBeforeHandling(
const mojom::ToolUseEvent& tool_use) const {
if (user_has_granted_permission_) {
return false;
}
// The search itself runs entirely on-device against the local embeddings
// index, so the query and the full history never leave the device. Only
// the matching pages' titles, URLs, and indexed passages are sent to the
// remote LLM as the tool's output. Ask the user to confirm before doing
// that. The user-facing wording is in
// `get_tool_permission_implications.tsx` so it goes through i18n; this
// C++ side only needs to surface a non-null challenge.
return mojom::PermissionChallenge::New(/*assessment=*/std::nullopt,
/*plan=*/std::nullopt);
}
void HistorySearchTool::UserPermissionGranted(const std::string& tool_use_id) {
user_has_granted_permission_ = true;
}
void HistorySearchTool::UseTool(const std::string& input_json,
UseToolCallback callback) {
auto input = base::JSONReader::ReadDict(input_json,
base::JSON_PARSE_CHROMIUM_EXTENSIONS);
if (!input.has_value()) {
std::move(callback).Run(
CreateContentBlocksForText("Error: failed to parse input JSON"), {});
return;
}
const std::string* query = input->FindString(kPropertyQuery);
if (!query || query->empty()) {
std::move(callback).Run(
CreateContentBlocksForText("Error: missing or empty 'query' field"),
{});
return;
}
// Match the chrome://history and omnibox gating: reject queries with fewer
// than `search_query_minimum_word_count` words (default 2). The service
// API itself does not enforce this, but very short queries produce noisy
// semantic results.
const int min_words = history_embeddings::GetFeatureParameters()
.search_query_minimum_word_count;
if (static_cast<int>(history_embeddings::CountWords(*query)) < min_words) {
std::move(callback).Run(
CreateContentBlocksForText(base::StrCat(
{"Error: 'query' must contain at least ",
base::NumberToString(min_words),
" words. Provide a longer natural-language description."})),
{});
return;
}
size_t count = kDefaultResultCount;
if (auto requested = input->FindInt(kPropertyCount)) {
count = static_cast<size_t>(
std::clamp(*requested, 1, static_cast<int>(kMaxResultCount)));
}
std::optional<base::Time> time_range_start;
if (auto days = input->FindInt(kPropertyTimeRangeStartDaysAgo);
days && *days > 0) {
time_range_start = base::Time::Now() - base::Days(*days);
}
bool include_all_passages =
input->FindBool(kPropertyIncludeAllPassages).value_or(true);
// Defense in depth: the tool is only published when the flag is on, but
// re-check here in case the flag flipped between advertisement and use.
// Tests that inject a fake search interface bypass this gate.
if (!search_for_testing_ &&
!history_embeddings::IsHistoryEmbeddingsEnabledForProfile(profile_)) {
std::move(callback).Run(
CreateContentBlocksForText(
"Error: history embeddings is not enabled for this profile"),
{});
return;
}
history_embeddings::HistoryEmbeddingsSearch* service =
search_for_testing_
? search_for_testing_.get()
: HistoryEmbeddingsServiceFactory::GetForProfile(profile_);
if (!service) {
std::move(callback).Run(
CreateContentBlocksForText(
"Error: history embeddings service is unavailable"),
{});
return;
}
// history_embeddings::SearchResultCallback is a RepeatingCallback. Adapt
// the OnceCallback by heap-allocating it and letting base::Owned() tie its
// lifetime to the bound closure; OnSearchResult guards against the
// theoretical case of more than one fire.
service->Search(
/*previous_search_result=*/nullptr, *query, time_range_start, count,
/*skip_answering=*/true,
base::BindRepeating(
&HistorySearchTool::OnSearchResult, weak_ptr_factory_.GetWeakPtr(),
base::Owned(std::make_unique<UseToolCallback>(std::move(callback))),
include_all_passages));
}
void HistorySearchTool::SetSearchForTesting(
history_embeddings::HistoryEmbeddingsSearch* search) {
search_for_testing_ = search;
}
void HistorySearchTool::OnSearchResult(
UseToolCallback* callback,
bool include_all_passages,
history_embeddings::SearchResult result) {
if (callback->is_null()) {
return;
}
std::string json = internal::BuildHistorySearchResultJson(
result.query, result, include_all_passages);
// Surface the URLs in our results as a visited-links artifact so the
// conversation UI permits anchors in the assistant's reply for them.
// The artifact has no visual rendering -- it's a sidechannel trust-list.
ToolArtifacts artifacts;
if (!result.scored_url_rows.empty()) {
base::ListValue links;
for (const auto& row : result.scored_url_rows) {
if (row.row.url().is_valid() &&
row.row.url().SchemeIs(url::kHttpsScheme)) {
links.Append(row.row.url().spec());
}
}
if (!links.empty()) {
std::string links_json;
base::JSONWriter::Write(links, &links_json);
artifacts.push_back(mojom::ToolArtifact::New(
/*id=*/std::nullopt, mojom::kVisitedLinksArtifactType,
std::move(links_json)));
}
}
std::move(*callback).Run(CreateContentBlocksForText(json),
std::move(artifacts));
}
} // namespace ai_chat
@@ -0,0 +1,88 @@
// Copyright (c) 2026 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.
#ifndef BRAVE_BROWSER_AI_CHAT_TOOLS_HISTORY_SEARCH_TOOL_H_
#define BRAVE_BROWSER_AI_CHAT_TOOLS_HISTORY_SEARCH_TOOL_H_
#include <optional>
#include <string>
#include <string_view>
#include <variant>
#include <vector>
#include "base/memory/raw_ptr.h"
#include "base/memory/weak_ptr.h"
#include "base/values.h"
#include "brave/components/ai_chat/core/browser/tools/tool.h"
#include "components/history_embeddings/core/history_embeddings_search.h"
class Profile;
namespace content {
class BrowserContext;
} // namespace content
namespace ai_chat {
namespace internal {
// Serializes a history embeddings SearchResult into the JSON payload
// returned by HistorySearchTool. Exposed for unit tests.
std::string BuildHistorySearchResultJson(
const std::string& query,
const history_embeddings::SearchResult& result,
bool include_all_passages);
} // namespace internal
// Exposes Brave's local history semantic search (history embeddings) as a
// tool callable by the AI Chat assistant. The tool runs the user's natural
// language query against the on-device embeddings index and returns the
// matching pages as a JSON-encoded TextContentBlock.
class HistorySearchTool : public Tool {
public:
explicit HistorySearchTool(content::BrowserContext* browser_context);
~HistorySearchTool() override;
HistorySearchTool(const HistorySearchTool&) = delete;
HistorySearchTool& operator=(const HistorySearchTool&) = delete;
// Tool:
std::string_view Name() const override;
std::string_view Description() const override;
std::optional<base::DictValue> InputProperties() const override;
std::optional<std::vector<std::string>> RequiredProperties() const override;
std::variant<bool, mojom::PermissionChallengePtr>
RequiresUserInteractionBeforeHandling(
const mojom::ToolUseEvent& tool_use) const override;
void UserPermissionGranted(const std::string& tool_use_id) override;
void UseTool(const std::string& input_json,
UseToolCallback callback) override;
// Tests inject a fake search interface to bypass the factory lookup and
// the per-profile enablement gate. The pointer must outlive any UseTool
// call made on this instance.
void SetSearchForTesting(history_embeddings::HistoryEmbeddingsSearch* search);
private:
// SearchResultCallback is a RepeatingCallback but with `skip_answering=true`
// it's expected to fire once. `callback` is heap-allocated and owned by the
// bound closure; the first invocation moves it out and runs the
// UseToolCallback, any further invocations are no-ops.
void OnSearchResult(UseToolCallback* callback,
bool include_all_passages,
history_embeddings::SearchResult result);
raw_ptr<Profile> profile_;
raw_ptr<history_embeddings::HistoryEmbeddingsSearch> search_for_testing_ =
nullptr;
bool user_has_granted_permission_ = false;
base::WeakPtrFactory<HistorySearchTool> weak_ptr_factory_{this};
};
} // namespace ai_chat
#endif // BRAVE_BROWSER_AI_CHAT_TOOLS_HISTORY_SEARCH_TOOL_H_
@@ -0,0 +1,301 @@
// Copyright (c) 2026 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.
#include "brave/browser/ai_chat/tools/history_search_tool.h"
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "base/functional/bind.h"
#include "base/i18n/time_formatting.h"
#include "base/json/json_writer.h"
#include "base/location.h"
#include "base/strings/strcat.h"
#include "base/task/sequenced_task_runner.h"
#include "base/test/run_until.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/test_future.h"
#include "base/test/values_test_util.h"
#include "base/time/time.h"
#include "base/values.h"
#include "brave/browser/ai_chat/ai_chat_conversation_ui_browsertest_base.h"
#include "brave/browser/ai_chat/browser_tool_provider.h"
#include "brave/components/ai_chat/core/browser/conversation_handler.h"
#include "brave/components/ai_chat/core/browser/engine/engine_consumer.h"
#include "brave/components/ai_chat/core/browser/engine/mock_engine_consumer.h"
#include "brave/components/ai_chat/core/browser/tools/tool_provider.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "components/history/core/browser/url_row.h"
#include "components/history_embeddings/core/history_embeddings_features.h"
#include "components/history_embeddings/core/history_embeddings_search.h"
#include "components/history_embeddings/core/vector_database.h"
#include "content/public/test/browser_test.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"
namespace ai_chat {
using ::testing::_;
using ::testing::AnyNumber;
using ::testing::Sequence;
namespace {
// Canned visit time for the row returned by the fake. The expected
// last_visit_time field in test JSON is derived from this same value via
// TimeFormatAsIso8601 so the two halves can't drift.
const base::Time& CannedVisitTime() {
static const base::Time time = base::Time::FromTimeT(1700000000);
return time;
}
// Minimal fake that records args and fires a canned result. The real
// HistoryEmbeddingsService also implements this interface, so the tool sees
// nothing unusual.
class FakeHistoryEmbeddingsSearch
: public history_embeddings::HistoryEmbeddingsSearch {
public:
history_embeddings::SearchResult Search(
history_embeddings::SearchResult* previous_search_result,
std::string query,
std::optional<base::Time> time_range_start,
size_t count,
bool skip_answering,
history_embeddings::SearchResultCallback callback) override {
last_query_ = query;
last_time_range_start_ = time_range_start;
last_count_ = count;
last_skip_answering_ = skip_answering;
history_embeddings::SearchResult result;
result.query = query;
history_embeddings::ScoredUrl scored_url(
/*url_id=*/1, /*visit_id=*/1, base::Time::Now(),
/*score=*/0.9f, /*word_match_score=*/0.0f);
history_embeddings::ScoredUrlRow row(scored_url);
row.row.set_url(GURL(canned_url_));
row.row.set_title(u"Canned title");
row.row.set_last_visit(CannedVisitTime());
result.scored_url_rows.push_back(std::move(row));
base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
FROM_HERE, base::BindOnce(callback, std::move(result)));
return history_embeddings::SearchResult();
}
std::string last_query_;
std::optional<base::Time> last_time_range_start_;
size_t last_count_ = 0;
bool last_skip_answering_ = false;
std::string canned_url_ = "https://example.com/canned";
};
} // namespace
class HistorySearchToolBrowserTest
: public AIChatConversationUIBrowserTestBase {
public:
HistorySearchToolBrowserTest() {
scoped_feature_list_.InitAndEnableFeature(
history_embeddings::kHistoryEmbeddings);
}
protected:
// Reroutes the conversation's production HistorySearchTool to a fake
// search service so that real HistoryEmbeddingsService calls are avoided.
void InjectFakeSearch(FakeHistoryEmbeddingsSearch* fake) {
auto* provider = static_cast<BrowserToolProvider*>(
conversation_handler_->GetFirstToolProviderForTesting());
ASSERT_TRUE(provider);
auto* tool = provider->GetHistorySearchToolForTesting();
ASSERT_TRUE(tool);
tool->SetSearchForTesting(fake);
}
// Waits for the tool's permission challenge to land in conversation
// history, then approves it so the tool can execute.
void WaitForAndApprovePermissionChallenge(
const std::string& tool_id,
base::Location location = base::Location::Current()) {
SCOPED_TRACE(location.ToString());
ASSERT_TRUE(base::test::RunUntil([this]() {
for (const auto& turn : conversation_handler_->GetConversationHistory()) {
if (!turn->events.has_value()) {
continue;
}
for (const auto& event : *turn->events) {
if (event->is_tool_use_event() &&
event->get_tool_use_event()->permission_challenge) {
return true;
}
}
}
return false;
}));
conversation_handler_->ProcessPermissionChallenge(tool_id, true);
}
base::test::ScopedFeatureList scoped_feature_list_;
};
// End-to-end: the engine emits a tool_use event referencing
// semantic_history_search; ConversationHandler dispatches to the real
// HistorySearchTool; the tool forwards the parsed arguments to the search
// service and returns the JSON-serialized result back into the conversation.
IN_PROC_BROWSER_TEST_F(HistorySearchToolBrowserTest,
ToolUseDispatchesSearchAndReturnsResult) {
CreateConversationWithMockEngine();
FakeHistoryEmbeddingsSearch fake;
InjectFakeSearch(&fake);
// After our tool returns, ConversationHandler will issue a follow-up
// GenerateAssistantResponse to feed the tool result back to the model.
// We don't care about that second call -- absorb it.
Sequence seq;
auto generate_future = SetupMockGenerateAssistantResponse(&seq);
EXPECT_CALL(*mock_engine_, GenerateAssistantResponse(_, _, _, _, _, _, _, _))
.Times(AnyNumber())
.InSequence(seq);
conversation_handler_->SubmitHumanConversationEntry(
"Find the article I read about kittens", std::nullopt);
auto callbacks = generate_future->Take();
base::DictValue args;
args.Set("query", "kittens article");
args.Set("count", 3);
args.Set("time_range_start_days_ago", 7);
args.Set("include_all_passages", false);
std::string args_json = *base::WriteJson(args);
callbacks.data_callback.Run(EngineConsumer::GenerationResultData(
mojom::ConversationEntryEvent::NewToolUseEvent(mojom::ToolUseEvent::New(
mojom::kSemanticHistorySearchToolName, "tool_id_1", args_json,
std::nullopt, std::nullopt, nullptr, false)),
std::nullopt));
std::move(callbacks.completed_callback)
.Run(base::ok(
EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
// The tool requires user permission before sending results to the LLM.
WaitForAndApprovePermissionChallenge("tool_id_1");
// Wait for ConversationHandler to dispatch to our tool, which forwards
// the query to the fake search service.
ASSERT_TRUE(
base::test::RunUntil([&fake]() { return !fake.last_query_.empty(); }));
EXPECT_EQ(fake.last_query_, "kittens article");
EXPECT_EQ(fake.last_count_, 3u);
EXPECT_TRUE(fake.last_skip_answering_)
<< "The tool must request the embeddings-only path -- answering is "
"handled by the Leo model itself.";
ASSERT_TRUE(fake.last_time_range_start_.has_value());
EXPECT_LT(base::Time::Now() - *fake.last_time_range_start_,
base::Days(7) + base::Minutes(5));
// The tool's output is plumbed back into the conversation as the
// tool_use_event's `output`. Wait for it to land before inspecting.
auto find_completed_tool_use = [this]() -> const mojom::ToolUseEvent* {
for (const auto& turn : conversation_handler_->GetConversationHistory()) {
if (!turn->events.has_value()) {
continue;
}
for (const auto& event : *turn->events) {
if (!event->is_tool_use_event()) {
continue;
}
const auto& tool_use = event->get_tool_use_event();
if (tool_use->tool_name == mojom::kSemanticHistorySearchToolName &&
tool_use->output.has_value() && !tool_use->output->empty()) {
return tool_use.get();
}
}
}
return nullptr;
};
ASSERT_TRUE(base::test::RunUntil(
[&find_completed_tool_use]() { return find_completed_tool_use(); }));
const mojom::ToolUseEvent* tool_use = find_completed_tool_use();
ASSERT_TRUE(tool_use);
const auto& blocks = *tool_use->output;
ASSERT_FALSE(blocks.empty());
ASSERT_TRUE(blocks[0]->is_text_content_block());
EXPECT_THAT(blocks[0]->get_text_content_block()->text,
base::test::IsJson(
base::StrCat({R"({
"query": "kittens article",
"results": [{
"title": "Canned title",
"url": "https://example.com/canned",
"last_visit_time": ")",
base::TimeFormatAsIso8601(CannedVisitTime()),
R"("
}]
})"})));
}
// When the engine sends malformed JSON args, the tool surfaces an error
// rather than crashing, and the conversation continues.
IN_PROC_BROWSER_TEST_F(HistorySearchToolBrowserTest,
MalformedArgsProducesErrorOutput) {
CreateConversationWithMockEngine();
FakeHistoryEmbeddingsSearch fake;
InjectFakeSearch(&fake);
// After our tool returns, ConversationHandler will issue a follow-up
// GenerateAssistantResponse to feed the tool result back to the model.
// We don't care about that second call -- absorb it.
Sequence seq;
auto generate_future = SetupMockGenerateAssistantResponse(&seq);
EXPECT_CALL(*mock_engine_, GenerateAssistantResponse(_, _, _, _, _, _, _, _))
.Times(AnyNumber())
.InSequence(seq);
conversation_handler_->SubmitHumanConversationEntry("anything", std::nullopt);
auto callbacks = generate_future->Take();
callbacks.data_callback.Run(EngineConsumer::GenerationResultData(
mojom::ConversationEntryEvent::NewToolUseEvent(mojom::ToolUseEvent::New(
mojom::kSemanticHistorySearchToolName, "tool_id_err", "{ not json",
std::nullopt, std::nullopt, nullptr, false)),
std::nullopt));
std::move(callbacks.completed_callback)
.Run(base::ok(
EngineConsumer::GenerationResultData(nullptr, std::nullopt)));
// The permission challenge fires before UseTool runs; approve it so the
// tool actually executes and surfaces its parse error.
WaitForAndApprovePermissionChallenge("tool_id_err");
// Wait for ConversationHandler to dispatch and our tool to surface its
// error output back into the conversation.
ASSERT_TRUE(base::test::RunUntil([this]() {
for (const auto& turn : conversation_handler_->GetConversationHistory()) {
if (!turn->events.has_value()) {
continue;
}
for (const auto& event : *turn->events) {
if (event->is_tool_use_event() &&
event->get_tool_use_event()->tool_name ==
mojom::kSemanticHistorySearchToolName &&
event->get_tool_use_event()->output.has_value() &&
!event->get_tool_use_event()->output->empty()) {
return true;
}
}
}
return false;
}));
EXPECT_TRUE(fake.last_query_.empty())
<< "Search service must not be called with unparseable args";
}
} // namespace ai_chat
@@ -0,0 +1,254 @@
// Copyright (c) 2026 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.
#include "brave/browser/ai_chat/tools/history_search_tool.h"
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
#include "base/strings/strcat.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/test_future.h"
#include "base/test/values_test_util.h"
#include "base/time/time.h"
#include "brave/components/ai_chat/core/common/mojom/common.mojom.h"
#include "chrome/test/base/testing_profile.h"
#include "components/history/core/browser/url_row.h"
#include "components/history_embeddings/content/history_embeddings_service.h"
#include "components/history_embeddings/core/history_embeddings_features.h"
#include "components/history_embeddings/core/history_embeddings_search.h"
#include "components/history_embeddings/core/vector_database.h"
#include "components/passage_embeddings/core/passage_embeddings_types.h"
#include "content/public/test/browser_task_environment.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"
namespace ai_chat {
namespace {
std::string ExtractText(const std::vector<mojom::ContentBlockPtr>& blocks) {
if (blocks.empty() || !blocks[0]->is_text_content_block()) {
return std::string();
}
return blocks[0]->get_text_content_block()->text;
}
std::string RunTool(HistorySearchTool* tool, const std::string& json) {
base::test::TestFuture<std::vector<mojom::ContentBlockPtr>,
std::vector<mojom::ToolArtifactPtr>>
future;
tool->UseTool(json, future.GetCallback());
return ExtractText(future.Get<std::vector<mojom::ContentBlockPtr>>());
}
history_embeddings::ScoredUrlRow MakeRow(const std::string& url,
const std::u16string& title) {
history_embeddings::ScoredUrl scored_url(/*url_id=*/1, /*visit_id=*/1,
base::Time::Now(),
/*score=*/0.9f,
/*word_match_score=*/0.0f);
history_embeddings::ScoredUrlRow row(scored_url);
row.row.set_url(GURL(url));
row.row.set_title(title);
row.row.set_last_visit(base::Time::FromTimeT(1700000000));
return row;
}
// Attaches a passage + matching embedding/score to a ScoredUrlRow. The
// embedding's data is irrelevant for serialization; only the score and word
// count are read by GetBestScoreIndices.
void AddPassage(history_embeddings::ScoredUrlRow& row,
const std::string& passage,
float score,
size_t word_count) {
row.url_data.passages.add_passages(passage);
// Embedding's ctor DCHECKs that the vector has unit magnitude; {1.0f} does.
row.url_data.passage_embeddings.emplace_back(
history_embeddings::PassageEmbedding{
passage_embeddings::Embedding(std::vector<float>{1.0f}), word_count});
row.scores.push_back(score);
}
} // namespace
class HistorySearchToolTest : public testing::Test {
protected:
HistorySearchToolTest()
: task_environment_(content::BrowserTaskEnvironment::IO_MAINLOOP) {}
void EnableFeature() {
feature_list_.InitAndEnableFeature(history_embeddings::kHistoryEmbeddings);
}
void DisableFeature() {
feature_list_.InitAndDisableFeature(history_embeddings::kHistoryEmbeddings);
}
std::unique_ptr<HistorySearchTool> CreateTool() {
return std::make_unique<HistorySearchTool>(&profile_);
}
content::BrowserTaskEnvironment task_environment_;
base::test::ScopedFeatureList feature_list_;
TestingProfile profile_;
};
TEST_F(HistorySearchToolTest, RequiresUserInteractionBeforeHandling) {
EnableFeature();
auto tool = CreateTool();
mojom::ToolUseEvent event(tool->Name().data(), "1", "{}", std::nullopt,
std::nullopt, nullptr, false);
// Before permission is granted: returns a PermissionChallenge so the
// user can confirm that history results may be sent to Brave AI. The
// user-facing wording lives in `get_tool_permission_implications.tsx`,
// so the challenge itself just needs to be a non-null sentinel here.
auto result = tool->RequiresUserInteractionBeforeHandling(event);
ASSERT_TRUE(std::holds_alternative<mojom::PermissionChallengePtr>(result));
EXPECT_TRUE(std::get<mojom::PermissionChallengePtr>(result));
// After permission is granted: no further interaction needed.
tool->UserPermissionGranted("1");
result = tool->RequiresUserInteractionBeforeHandling(event);
ASSERT_TRUE(std::holds_alternative<bool>(result));
EXPECT_FALSE(std::get<bool>(result));
}
TEST_F(HistorySearchToolTest, InvalidInputProducesError) {
EnableFeature();
auto tool = CreateTool();
struct {
std::string_view name;
std::string_view input_json;
std::string_view expected_error_substring;
} kCases[] = {
{"malformed json", "{ not json", "failed to parse"},
{"missing query", R"({})", "missing or empty 'query'"},
{"empty query", R"({"query": ""})", "missing or empty 'query'"},
// Default search_query_minimum_word_count is 2.
{"single-word query", R"({"query": "kittens"})", "at least 2 words"},
};
for (const auto& c : kCases) {
SCOPED_TRACE(c.name);
EXPECT_THAT(RunTool(tool.get(), std::string(c.input_json)),
testing::HasSubstr(std::string(c.expected_error_substring)));
}
}
TEST_F(HistorySearchToolTest, FeatureDisabledReturnsError) {
DisableFeature();
auto tool = CreateTool();
EXPECT_THAT(RunTool(tool.get(), R"({"query": "anything at all"})"),
testing::HasSubstr("not enabled"));
}
TEST_F(HistorySearchToolTest, InputPropertiesDeclareQueryRequired) {
auto tool = CreateTool();
auto required = tool->RequiredProperties();
ASSERT_TRUE(required.has_value());
EXPECT_THAT(*required, testing::Contains("query"));
}
TEST_F(HistorySearchToolTest, InputPropertiesDeclareOptionalProperties) {
auto tool = CreateTool();
auto props = tool->InputProperties();
ASSERT_TRUE(props.has_value());
EXPECT_TRUE(props->Find("query"));
EXPECT_TRUE(props->Find("count"));
EXPECT_TRUE(props->Find("time_range_start_days_ago"));
}
// The visit time set by MakeRow is base::Time::FromTimeT(1700000000),
// which TimeFormatAsIso8601 renders as this exact string.
constexpr char kExpectedLastVisitTime[] = "2023-11-14T22:13:20.000Z";
TEST(HistorySearchToolJsonTest, EmptyResultProducesEmptyResultsArray) {
history_embeddings::SearchResult result;
std::string json = internal::BuildHistorySearchResultJson(
"kittens", result, /*include_all_passages=*/false);
EXPECT_THAT(json,
base::test::IsJson(R"({"query": "kittens", "results": []})"));
}
TEST(HistorySearchToolJsonTest, ResultsAreSerialized) {
history_embeddings::SearchResult result;
result.scored_url_rows.push_back(
MakeRow("https://example.com/article", u"Article title"));
std::string json = internal::BuildHistorySearchResultJson(
"article", result, /*include_all_passages=*/false);
// No passages were attached to the row, so neither 'passages' nor
// 'snippet' is emitted -- otherwise BuildHistorySearchResultJson would
// trigger GetBestPassage()'s CHECK.
EXPECT_THAT(json, base::test::IsJson(base::StrCat({R"({
"query": "article",
"results": [{
"title": "Article title",
"url": "https://example.com/article",
"last_visit_time": ")",
kExpectedLastVisitTime,
R"("
}]
})"})));
}
TEST(HistorySearchToolJsonTest, IncludeAllPassagesOrdersByScore) {
history_embeddings::SearchResult result;
auto row = MakeRow("https://example.com/article", u"Article title");
AddPassage(row, "low score passage", /*score=*/0.1f, /*word_count=*/3);
AddPassage(row, "high score passage", /*score=*/0.9f, /*word_count=*/3);
AddPassage(row, "mid score passage", /*score=*/0.5f, /*word_count=*/3);
result.scored_url_rows.push_back(std::move(row));
std::string json = internal::BuildHistorySearchResultJson(
"article", result, /*include_all_passages=*/true);
EXPECT_THAT(json, base::test::IsJson(base::StrCat({R"({
"query": "article",
"results": [{
"title": "Article title",
"url": "https://example.com/article",
"last_visit_time": ")",
kExpectedLastVisitTime,
R"(",
"passages": [
"high score passage",
"mid score passage",
"low score passage"
]
}]
})"})));
}
TEST(HistorySearchToolJsonTest, BestPassageSnippetWhenNotIncludingAll) {
history_embeddings::SearchResult result;
auto row = MakeRow("https://example.com/article", u"Article title");
AddPassage(row, "low score passage", /*score=*/0.1f, /*word_count=*/3);
AddPassage(row, "high score passage", /*score=*/0.9f, /*word_count=*/3);
result.scored_url_rows.push_back(std::move(row));
std::string json = internal::BuildHistorySearchResultJson(
"article", result, /*include_all_passages=*/false);
EXPECT_THAT(json, base::test::IsJson(base::StrCat({R"({
"query": "article",
"results": [{
"title": "Article title",
"url": "https://example.com/article",
"last_visit_time": ")",
kExpectedLastVisitTime,
R"(",
"snippet": "high score passage"
}]
})"})));
}
} // namespace ai_chat
@@ -20,6 +20,7 @@ const string kUserChoiceToolName = "user_choice_tool";
const string kAssistantDetailStorageToolName = "assistant_detail_storage";
const string kMemoryStorageToolName = "memory_storage_tool";
const string kCodeExecutionToolName = "code_execution_tool";
const string kSemanticHistorySearchToolName = "semantic_history_search";
[EnableIf=enable_ai_chat_tab_management_tool]
const string kTabManagementToolName = "tab_management";
@@ -45,6 +46,13 @@ const string kWebContextSearchToolName = "web_context_search";
// Tool artifact types
const string kLineChartArtifactType = "line_chart";
// Carries a JSON array of HTTPS URL strings that the tool's output
// references -- typically pages from the user's local browsing history
// that match the tool's query. The conversation UI uses this to permit
// the assistant's reply to render those URLs as anchors. Not rendered
// visually.
const string kVisitedLinksArtifactType = "visited_links";
// Maximum length for memory and customization records.
const int32 kMaxMemoryRecordLength = 512;
@@ -35,6 +35,19 @@ function getHistoryToolNameLabel(toolInput: any) {
return getLocale(S.CHAT_UI_TOOL_LABEL_NAVIGATE_WEB_PAGE)
}
/**
* See history_search_tool.cc
* @param toolInput Expects { "query": string } but is not guaranteed.
*/
function getSemanticHistorySearchToolNameLabel(toolInput: any) {
if (typeof toolInput?.query === 'string' && toolInput.query.length > 0) {
return formatLocale(S.CHAT_UI_TOOL_LABEL_SEMANTIC_HISTORY_SEARCH, {
$1: toolInput.query as string,
})
}
return getLocale(S.CHAT_UI_TOOL_LABEL_SEMANTIC_HISTORY_SEARCH_GENERIC)
}
function getSearchToolNameLabel(toolInput: any) {
// toolInput is parsed (possibly malformed) JSON, so it may be undefined and
// its `query` field — confusingly named — may be missing or not actually be
@@ -78,6 +91,8 @@ export function getToolLabel(toolName: string, toolInput: any) {
return getLocale(S.CHAT_UI_TOOL_LABEL_WAIT)
case Mojom.CODE_EXECUTION_TOOL_NAME:
return getLocale(S.CHAT_UI_TOOL_LABEL_CODE_EXECUTION)
case Mojom.SEMANTIC_HISTORY_SEARCH_TOOL_NAME:
return getSemanticHistorySearchToolNameLabel(toolInput)
// <if expr="enable_ai_chat_tab_management_tool">
case Mojom.TAB_MANAGEMENT_TOOL_NAME:
return getLocale(S.CHAT_UI_TOOL_LABEL_TAB_MANAGEMENT)
@@ -3,14 +3,23 @@
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.
import {
getLocale,
// <if expr="enable_ai_chat_tab_management_tool">
formatLocale,
// </if>
} from '$web-common/locale'
import * as Mojom from '../../../common/mojom'
// <if expr="enable_ai_chat_tab_management_tool">
import * as React from 'react'
import { formatLocale } from '$web-common/locale'
import * as Mojom from '../../../common/mojom'
// </if>
export function getToolPermissionImplications(toolName: string) {
switch (toolName) {
case Mojom.SEMANTIC_HISTORY_SEARCH_TOOL_NAME:
return getLocale(
S.CHAT_UI_TOOL_SEMANTIC_HISTORY_SEARCH_PERMISSION_IMPLICATIONS,
)
// <if expr="enable_ai_chat_tab_management_tool">
case Mojom.TAB_MANAGEMENT_TOOL_NAME:
return formatLocale(
@@ -67,6 +67,7 @@ export const _AssistantTask = {
assistantEntries={taskConversationEntries}
isActiveTask={args.isActiveTask}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>
)
@@ -65,6 +65,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -101,6 +102,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -127,6 +129,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -161,6 +164,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -219,6 +223,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -262,6 +267,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -327,6 +333,7 @@ describe('AssistantTask', () => {
assistantEntries={entriesWithInlineSearch}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -357,6 +364,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={true}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -384,6 +392,7 @@ describe('AssistantTask', () => {
assistantEntries={mockAssistantEntries}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -455,6 +464,7 @@ describe('AssistantTask web sources', () => {
assistantEntries={entriesWithSourcesInEarlierEntry}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</MockContext>,
)
@@ -473,6 +483,7 @@ describe('AssistantTask web sources', () => {
assistantEntries={entriesWithSourcesInEarlierEntry}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</ProgressBubbleContextProvider>
</MockContext>,
@@ -507,6 +518,7 @@ describe('AssistantTask web sources', () => {
assistantEntries={entriesWithSourcesInEarlierEntry}
isActiveTask={false}
isLeoModel={true}
allowedLinks={[]}
/>
</ProgressBubbleContextProvider>
</MockContext>,
@@ -26,6 +26,11 @@ interface Props {
// Passes to AssistantResponse
isLeoModel: boolean
// URLs the assistant's reply is allowed to render as anchors. Computed
// once per group by the caller so both the AssistantResponse and
// AssistantTask render paths see the same list.
allowedLinks: string[]
}
interface TabProps {
@@ -193,7 +198,7 @@ function Progress(props: Props & TabProps) {
]}
isEntryInteractivityAllowed={false}
isEntryInProgress={props.isGenerating}
allowedLinks={props.taskData.allowedLinks}
allowedLinks={props.allowedLinks}
isLeoModel={props.isLeoModel}
toolArtifacts={props.toolArtifacts}
/>
@@ -269,7 +274,7 @@ function Steps(props: Props & TabProps) {
events={events}
isEntryInteractivityAllowed={isRunnable}
isEntryInProgress={isActive}
allowedLinks={props.taskData.allowedLinks}
allowedLinks={props.allowedLinks}
isLeoModel={props.isLeoModel}
/>
</div>
@@ -9,7 +9,6 @@ import {
createConversationTurnWithDefaults,
getCompletionEvent,
getToolUseEvent,
getWebSourcesEvent,
} from '../../../common/test_data_utils'
import useExtractTaskData from './use_extract_task_data'
@@ -115,109 +114,4 @@ describe('useExtractTaskData', () => {
)
})
})
describe('extracting allowed links', () => {
it('should extract links from sources events', () => {
const assistantEntries: Mojom.ConversationTurn[] = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getWebSourcesEvent([
{
url: { url: 'https://example1.com' },
title: 'Example 1',
faviconUrl: { url: 'https://example1.com/favicon.ico' },
},
{
url: { url: 'https://example2.com' },
title: 'Example 2',
faviconUrl: { url: 'https://example2.com/favicon.ico' },
},
]),
getCompletionEvent('Task'),
],
}),
]
const { result } = renderHook(() => useExtractTaskData(assistantEntries))
expect(result.current.allowedLinks).toHaveLength(2)
expect(result.current.allowedLinks).toContain('https://example1.com')
expect(result.current.allowedLinks).toContain('https://example2.com')
})
it('should accumulate links from multiple sources events', () => {
const assistantEntries: Mojom.ConversationTurn[] = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getWebSourcesEvent([
{
url: { url: 'https://example1.com' },
title: 'Example 1',
faviconUrl: { url: 'https://example1.com/favicon.ico' },
},
{
url: { url: 'https://example2.com' },
title: 'Example 2',
faviconUrl: { url: 'https://example2.com/favicon.ico' },
},
]),
getCompletionEvent('First task'),
],
}),
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getWebSourcesEvent([
{
url: { url: 'https://example2.com' },
title: 'Duplicate Example 2',
faviconUrl: { url: 'https://example2.com/favicon.ico' },
},
{
url: { url: 'https://example3.com' },
title: 'Example 3',
faviconUrl: { url: 'https://example3.com/favicon.ico' },
},
]),
getCompletionEvent('Second task'),
],
}),
]
const { result } = renderHook(() => useExtractTaskData(assistantEntries))
expect(result.current.allowedLinks).toHaveLength(3) // ignores duplicate
expect(result.current.allowedLinks).toContain('https://example1.com')
expect(result.current.allowedLinks).toContain('https://example2.com')
expect(result.current.allowedLinks).toContain('https://example3.com')
})
it('should handle entries with no sources events', () => {
const assistantEntries: Mojom.ConversationTurn[] = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [getCompletionEvent('Task with no sources')],
}),
]
const { result } = renderHook(() => useExtractTaskData(assistantEntries))
expect(result.current.allowedLinks).toHaveLength(0)
})
it('should handle entries with no events', () => {
const assistantEntries: Mojom.ConversationTurn[] = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: undefined,
}),
]
const { result } = renderHook(() => useExtractTaskData(assistantEntries))
expect(result.current.allowedLinks).toHaveLength(0)
})
})
})
@@ -12,9 +12,6 @@ export default function useExtractTaskData(
return React.useMemo(() => {
// Individual tasks, split by CompletionEvent
const taskItems: Mojom.ConversationEntryEvent[][] = []
// All completion events are allowed the links provided by the whole response
// group.
const allowedLinks = new Set<string>()
for (const event of assistantEntryGroup.flatMap(
(entry) => entry.events ?? [],
@@ -28,18 +25,11 @@ export default function useExtractTaskData(
}
// Add any other event types to the last task item
taskItems.at(-1)?.push(event)
// Additionally collate the allowed links when a sources event is found
if (event.sourcesEvent) {
for (const source of event.sourcesEvent.sources) {
allowedLinks.add(source.url.url)
}
}
}
}
return {
taskItems,
allowedLinks: Array.from(allowedLinks),
}
}, [assistantEntryGroup])
}
@@ -5,6 +5,7 @@
import * as Mojom from '../../../common/mojom'
import {
getGroupAllowedLinks,
getReasoningText,
removeReasoning,
removeCitationsWithMissingLinks,
@@ -17,6 +18,7 @@ import {
createConversationTurnWithDefaults,
getToolUseEvent,
getCompletionEvent,
getWebSourcesEvent,
} from '../../../common/test_data_utils'
describe('groupConversationEntries', () => {
@@ -475,3 +477,164 @@ describe('isAssistantGroupTask', () => {
expect(isAssistantGroupTask(group)).toBe(false)
})
})
describe('getGroupAllowedLinks', () => {
const makeArtifact = (contentJson: string): Mojom.ToolArtifact => ({
id: null,
type: Mojom.VISITED_LINKS_ARTIFACT_TYPE,
contentJson,
})
const turnWithArtifacts = (
artifacts: Mojom.ToolArtifact[],
): Mojom.ConversationTurn =>
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getToolUseEvent({
toolName: Mojom.SEMANTIC_HISTORY_SEARCH_TOOL_NAME,
id: '1',
argumentsJson: '{}',
output: [],
artifacts,
}),
],
})
it('returns URL strings from a visited_links artifact', () => {
const group = [
turnWithArtifacts([
makeArtifact(
JSON.stringify(['https://a.com/one', 'https://a.com/two']),
),
]),
]
expect(getGroupAllowedLinks(group)).toEqual([
'https://a.com/one',
'https://a.com/two',
])
})
it('flattens across all entries in the group', () => {
const group = [
turnWithArtifacts([makeArtifact(JSON.stringify(['https://a.com']))]),
turnWithArtifacts([makeArtifact(JSON.stringify(['https://b.com']))]),
]
expect(getGroupAllowedLinks(group)).toEqual([
'https://a.com',
'https://b.com',
])
})
it('ignores artifacts of other types', () => {
const group = [
turnWithArtifacts([
{ id: null, type: 'line_chart', contentJson: '{"data": []}' },
makeArtifact(JSON.stringify(['https://kept.com'])),
]),
]
expect(getGroupAllowedLinks(group)).toEqual(['https://kept.com'])
})
it('filters out malformed JSON, non-array, and non-string entries', () => {
const group = [
turnWithArtifacts([
makeArtifact('not valid json'),
makeArtifact(JSON.stringify({ not: 'an array' })),
makeArtifact(JSON.stringify(['https://example.com/ok', 42, null])),
]),
]
expect(getGroupAllowedLinks(group)).toEqual(['https://example.com/ok'])
})
it('returns an empty array when no tool_use events exist', () => {
const group = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [getCompletionEvent('plain response')],
}),
]
expect(getGroupAllowedLinks(group)).toEqual([])
})
it('uses the latest edit of each entry', () => {
const original = turnWithArtifacts([
makeArtifact(JSON.stringify(['https://original.com'])),
])
original.edits = [
turnWithArtifacts([makeArtifact(JSON.stringify(['https://edited.com']))]),
]
expect(getGroupAllowedLinks([original])).toEqual(['https://edited.com'])
})
it('collects URLs from sourcesEvent across all entries in the group', () => {
const group = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getWebSourcesEvent([
{
url: { url: 'https://a.com' },
title: 'A',
faviconUrl: { url: 'https://a.com/favicon.ico' },
},
]),
],
}),
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getWebSourcesEvent([
{
url: { url: 'https://b.com' },
title: 'B',
faviconUrl: { url: 'https://b.com/favicon.ico' },
},
]),
],
}),
]
expect(getGroupAllowedLinks(group)).toEqual([
'https://a.com',
'https://b.com',
])
})
it('unions sourcesEvent and visited_links URLs and dedupes', () => {
const group = [
createConversationTurnWithDefaults({
characterType: Mojom.CharacterType.ASSISTANT,
events: [
getToolUseEvent({
toolName: Mojom.SEMANTIC_HISTORY_SEARCH_TOOL_NAME,
id: '1',
argumentsJson: '{}',
output: [],
artifacts: [
makeArtifact(
JSON.stringify(['https://history.example', 'https://dup.com']),
),
],
}),
getWebSourcesEvent([
{
url: { url: 'https://search.example' },
title: 'Search',
faviconUrl: { url: 'https://search.example/favicon.ico' },
},
{
url: { url: 'https://dup.com' },
title: 'Dup',
faviconUrl: { url: 'https://dup.com/favicon.ico' },
},
]),
],
}),
]
expect(getGroupAllowedLinks(group)).toEqual([
'https://history.example',
'https://dup.com',
'https://search.example',
])
})
})
@@ -279,3 +279,58 @@ export function getToolArtifacts(
return [...artifactsWithoutId, ...artifactsById.values()]
}
/**
* Collects every URL that should be permitted as an anchor in the assistant
* replies for this group, deduped. Combines:
* - Web search citations from `sourcesEvent` (the "Sources" panel URLs).
* - HTTPS URLs from `visited_links` artifacts on tool_use events --
* client-side tools (e.g. semantic history search) emit these as a
* sidechannel trust-list so the assistant's reply can render the tool's
* URLs as anchors. Bad JSON or non-string array entries are skipped.
* Flattening across the whole group is required because a client-side tool
* call lives in a separate assistant entry from the follow-up response that
* references the tool's URLs.
*/
function parseVisitedLinksArtifact(artifact: Mojom.ToolArtifact): string[] {
if (artifact.type !== Mojom.VISITED_LINKS_ARTIFACT_TYPE) {
return []
}
try {
const parsed: unknown = JSON.parse(artifact.contentJson)
return Array.isArray(parsed)
? parsed.filter((u): u is string => typeof u === 'string')
: []
} catch {
return []
}
}
function collectLinksFromEvent(
event: Mojom.ConversationEntryEvent,
links: Set<string>,
) {
if (event.sourcesEvent) {
for (const source of event.sourcesEvent.sources) {
links.add(source.url.url)
}
}
for (const a of event.toolUseEvent?.artifacts ?? []) {
for (const url of parseVisitedLinksArtifact(a)) {
links.add(url)
}
}
}
export function getGroupAllowedLinks(
group: Mojom.ConversationTurn[],
): string[] {
const links = new Set<string>()
for (const entry of group) {
const events = (entry.edits?.at(-1) ?? entry).events ?? []
for (const event of events) {
collectLinksFromEvent(event, links)
}
}
return Array.from(links)
}
@@ -112,29 +112,39 @@ describe('ConversationEntries allowedLinks per response', () => {
])
})
it('passes correct allowedLinks for a combined AssistantResponse group', () => {
render(
<MockContext
overrides={mockOverrides}
initialState={{
conversationHistory: [
humanTurn1,
assistantTurn1,
assistantTurn2,
] as any,
}}
>
<ConversationEntries />
</MockContext>,
)
expect(assistantResponseMock).toHaveBeenCalledTimes(2)
expect(assistantResponseMock.mock.calls[0][0]?.allowedLinks).toEqual([
'https://a.com',
])
expect(assistantResponseMock.mock.calls[1][0]?.allowedLinks).toEqual([
'https://b.com',
])
})
it(
'passes the same group-wide allowedLinks to every AssistantResponse '
+ 'in a combined group',
() => {
// A client-side tool call lives in a separate assistant entry from the
// follow-up response that references its URLs, so allowedLinks is
// computed per group (not per entry) and every entry in the group sees
// the union.
render(
<MockContext
overrides={mockOverrides}
initialState={{
conversationHistory: [
humanTurn1,
assistantTurn1,
assistantTurn2,
] as any,
}}
>
<ConversationEntries />
</MockContext>,
)
expect(assistantResponseMock).toHaveBeenCalledTimes(2)
expect(assistantResponseMock.mock.calls[0][0]?.allowedLinks).toEqual([
'https://a.com',
'https://b.com',
])
expect(assistantResponseMock.mock.calls[1][0]?.allowedLinks).toEqual([
'https://a.com',
'https://b.com',
])
},
)
})
describe('conversation entries', () => {
@@ -30,6 +30,7 @@ import AssistantResponse from '../assistant_response'
import EditInput from '../edit_input'
import EditIndicator from '../edit_indicator'
import {
getGroupAllowedLinks,
getReasoningText,
getToolArtifacts,
groupConversationEntries,
@@ -312,6 +313,13 @@ function ConversationEntries(props: { scrollToBottom: () => void }) {
? getToolArtifacts(group)
: null
// Computed once per group and passed to both AssistantTask and
// AssistantResponse so all entries in the group share the same set of
// anchor-permitted URLs. Required because a client-side tool call lives
// in a separate assistant entry from the follow-up response that
// references the tool's URLs.
const allowedLinks = getGroupAllowedLinks(group)
return (
<div key={firstEntryEdit.uuid || entryNumber}>
<div
@@ -327,6 +335,7 @@ function ConversationEntries(props: { scrollToBottom: () => void }) {
assistantEntries={group}
isActiveTask={isActiveGroup}
isLeoModel={conversationContext.isLeoModel}
allowedLinks={allowedLinks}
/>
)}
{!groupIsTask
@@ -336,13 +345,6 @@ function ConversationEntries(props: { scrollToBottom: () => void }) {
const isActiveEntryInActiveGroup =
isActiveGroup && i === group.length - 1
const currentEntryEdit = entry.edits?.at(-1) ?? entry
const allowedLinksForEntry: string[] =
currentEntryEdit.events?.flatMap(
(event) =>
event.sourcesEvent?.sources?.map(
(source) => source.url.url,
) || [],
) || []
const entryText = getCompletion(currentEntryEdit)
const hasReasoning = entryText.includes('<think>')
@@ -368,7 +370,7 @@ function ConversationEntries(props: { scrollToBottom: () => void }) {
isActiveEntryInActiveGroup
}
isEntryInProgress={isEntryInProgress}
allowedLinks={allowedLinksForEntry}
allowedLinks={allowedLinks}
isLeoModel={conversationContext.isLeoModel}
toolArtifacts={
i === group.length - 1 ? toolArtifacts : null
@@ -66,13 +66,11 @@ export default function ToolPermissionChallenge(props: Props) {
</>
)}
{toolPermissionImplications && <p>{toolPermissionImplications}</p>}
{props.toolUseEvent.permissionChallenge?.plan && (
<>
{toolPermissionImplications && <p>{toolPermissionImplications}</p>}
<p className={styles.assessment}>
{props.toolUseEvent.permissionChallenge?.plan}
</p>
</>
<p className={styles.assessment}>
{props.toolUseEvent.permissionChallenge.plan}
</p>
)}
<ConversationAreaButton
@@ -933,6 +933,15 @@
<message name="IDS_CHAT_UI_TOOL_LABEL_CODE_EXECUTION" desc="Title for code execution tool event" formatter_data="webui=AiChat">
Executing code in sandbox to help with your query
</message>
<message name="IDS_CHAT_UI_TOOL_LABEL_SEMANTIC_HISTORY_SEARCH" desc="Label shown while the assistant is searching the user's local browsing history for a query. $1 is the natural-language search query." formatter_data="webui=AiChat">
Searching your browser history for "<ph name="QUERY">$1<ex>kittens article</ex></ph>"
</message>
<message name="IDS_CHAT_UI_TOOL_LABEL_SEMANTIC_HISTORY_SEARCH_GENERIC" desc="Label shown while the assistant is searching the user's local browsing history when the query is not yet known." formatter_data="webui=AiChat">
Searching your browser history
</message>
<message name="IDS_CHAT_UI_TOOL_SEMANTIC_HISTORY_SEARCH_PERMISSION_IMPLICATIONS" desc="Explanation shown in the permission prompt when the assistant wants to search the user's local browsing history. Calls out that the search itself is on-device and only the matched pages flow to Brave AI." formatter_data="webui=AiChat">
Brave AI wants to search your browsing history. The search runs locally on your device, so your full history isn't sent anywhere. Only the matching pages' titles, URLs, and indexed text passages will be sent to Brave AI as part of this conversation.
</message>
<message name="IDS_CHAT_UI_CODE_EXECUTION_CODE_LABEL" desc="Label for code section in code execution tool" formatter_data="webui=AiChat">
Code:
</message>