Discharge7443/brave-core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[AI Chat] Add support for code execution plugins (#33456)

Browse Source 
* [AI Chat] Add support for code execution plugins

* Address PR feedback

* Update code after Chromium upgrade
This commit is contained in:
Darnell Andries
2026-02-26 20:51:44 -08:00
committed by GitHub
parent 2b14c4bb3c
commit 60235f3ea4
8 changed files with 315 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
browser/ai_chat/BUILD.gn
+2 -2
View File
@@ -23,12 +23,12 @@ static_library("ai_chat") {
"browser_tool_provider.h",
"browser_tool_provider_factory.cc",
"browser_tool_provider_factory.h",
"code_execution_tool.cc",
"code_execution_tool.h",
"tab_data_web_contents_observer.cc",
"tab_data_web_contents_observer.h",
"tab_tracker_service_factory.cc",
"tab_tracker_service_factory.h",
"tools/code_execution_tool.cc",
"tools/code_execution_tool.h",
"upload_file_helper.cc",
"upload_file_helper.h",
]
browser/ai_chat/browser_tool_provider.cc
+1 -1
View File
@@ -10,7 +10,7 @@
#include "base/feature_list.h"
#include "base/memory/weak_ptr.h"
#include "brave/browser/ai_chat/code_execution_tool.h"
#include "brave/browser/ai_chat/tools/code_execution_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"
browser/ai_chat/code_execution_tool_browsertest.cc
+93 -6
View File
@@ -3,10 +3,12 @@
* 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/code_execution_tool.h"
#include "brave/browser/ai_chat/tools/code_execution_tool.h"
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <vector>
#include "base/json/json_writer.h"
@@ -14,6 +16,7 @@
#include "base/strings/strcat.h"
#include "base/test/bind.h"
#include "base/values.h"
#include "brave/components/ai_chat/core/browser/tools/code_plugin.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/browser/ui/browser.h"
@@ -27,6 +30,40 @@ using testing::HasSubstr;
namespace ai_chat {
namespace {
// A minimal CodePlugin that exposes a global `mockPlugin` object with a
// `getValue()` method and handles "mock" artifacts.
class MockCodePlugin : public CodePlugin {
public:
// Allow tests to override ValidateArtifact behavior.
void SetValidationError(std::optional<std::string> error) {
validation_error_ = std::move(error);
}
std::string_view Description() const override {
return "Mock plugin for testing.";
}
std::string_view InclusionKeyword() const override { return "mockPlugin"; }
std::string_view SetupScript() const override {
return "const mockPlugin = { getValue: () => 'mock_value' };";
}
std::string_view ArtifactType() const override { return "mock"; }
std::optional<std::string> ValidateArtifact(
const base::Value& artifact_value) const override {
return validation_error_;
}
private:
std::optional<std::string> validation_error_;
};
} // namespace
class AIChatCodeExecutionToolBrowserTest : public InProcessBrowserTest {
public:
AIChatCodeExecutionToolBrowserTest() = default;
@@ -51,27 +88,37 @@ class AIChatCodeExecutionToolBrowserTest : public InProcessBrowserTest {
return http_server_.GetURL("/test").spec();
}
void ExecuteCodeRaw(const std::string& input_json, std::string* output) {
void ExecuteCodeRaw(
const std::string& input_json,
std::string* output,
std::vector<mojom::ToolArtifactPtr>* artifacts = nullptr) {
base::RunLoop run_loop;
tool_->UseTool(
input_json,
base::BindLambdaForTesting(
[&run_loop, output](std::vector<mojom::ContentBlockPtr> result,
std::vector<mojom::ToolArtifactPtr>) {
[&run_loop, output, artifacts](
std::vector<mojom::ContentBlockPtr> result,
std::vector<mojom::ToolArtifactPtr> result_artifacts) {
ASSERT_FALSE(result.empty());
ASSERT_TRUE(result[0]->is_text_content_block());
*output = result[0]->get_text_content_block()->text;
if (artifacts) {
*artifacts = std::move(result_artifacts);
}
run_loop.Quit();
}));
run_loop.Run();
}
void ExecuteCode(const std::string& script, std::string* output) {
void ExecuteCode(const std::string& script,
std::string* output,
std::vector<mojom::ToolArtifactPtr>* artifacts = nullptr) {
base::DictValue input;
input.Set("script", script);
std::string input_json;
base::JSONWriter::Write(input, &input_json);
ExecuteCodeRaw(input_json, output);
ExecuteCodeRaw(input_json, output, artifacts);
}
protected:
@@ -211,4 +258,44 @@ IN_PROC_BROWSER_TEST_F(AIChatCodeExecutionToolBrowserTest,
EXPECT_EQ(output, "0.3");
}
// Plugin setup script is injected when the keyword is present, the artifact is
// returned, and console output is captured correctly.
IN_PROC_BROWSER_TEST_F(AIChatCodeExecutionToolBrowserTest, PluginReturnsValue) {
tool_->AddCodePluginForTesting(std::make_unique<MockCodePlugin>());
std::string output;
std::vector<mojom::ToolArtifactPtr> artifacts;
ExecuteCode(
R"(
codeExecArtifacts.push({ type: 'mock', content: mockPlugin.getValue() });
console.log('done');
)",
&output, &artifacts);
EXPECT_EQ(output, "done");
ASSERT_EQ(artifacts.size(), 1u);
EXPECT_EQ(artifacts[0]->type, "mock");
EXPECT_EQ(artifacts[0]->content_json, "\"mock_value\"");
}
// A failed ValidateArtifact call surfaces as an error in the output.
IN_PROC_BROWSER_TEST_F(AIChatCodeExecutionToolBrowserTest,
PluginArtifactValidationError) {
auto plugin = std::make_unique<MockCodePlugin>();
plugin->SetValidationError("content must be a number");
tool_->AddCodePluginForTesting(std::move(plugin));
std::string output;
std::vector<mojom::ToolArtifactPtr> artifacts;
ExecuteCode(
R"(
codeExecArtifacts.push({ type: 'mock', content: mockPlugin.getValue() });
console.log('done');
)",
&output, &artifacts);
EXPECT_THAT(output, HasSubstr("Error: content must be a number"));
EXPECT_TRUE(artifacts.empty());
}
} // namespace ai_chat
browser/ai_chat/code_execution_tool.cc → browser/ai_chat/tools/code_execution_tool.cc
+141 -44
View File
@@ -3,11 +3,12 @@
// 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/code_execution_tool.h"
#include "brave/browser/ai_chat/tools/code_execution_tool.h"
#include <utility>
#include "base/json/json_reader.h"
#include "base/json/json_writer.h"
#include "base/strings/strcat.h"
#include "base/strings/string_util.h"
#include "base/strings/utf_string_conversions.h"
@@ -37,24 +38,16 @@ namespace {
constexpr base::TimeDelta kExecutionTimeLimit = base::Seconds(10);
constexpr char kScriptProperty[] = "script";
std::string WrapScript(const std::string& script) {
auto bignumber_js =
ui::ResourceBundle::GetSharedInstance().LoadDataResourceString(
IDR_AI_CHAT_BIGNUMBER_JS);
return base::StrCat({"(async function() { ", bignumber_js, " try { ", script,
" } catch (error) { console.error(error.toString()); } "
"return true; })()"});
}
constexpr char kArtifactTypeKey[] = "type";
constexpr char kArtifactContentKey[] = "content";
} // namespace
CodeExecutionTool::CodeExecutionRequest::CodeExecutionRequest(
Profile* profile,
const std::string& script,
std::string script,
base::TimeDelta execution_time_limit)
: content::WebContentsObserver(nullptr), wrapped_js_(WrapScript(script)) {
: content::WebContentsObserver(nullptr), script_(std::move(script)) {
auto otr_profile_id = Profile::OTRProfileID::AIChatCodeExecutionID();
auto* otr_profile = profile->GetOffTheRecordProfile(
otr_profile_id, /*create_if_needed=*/true);
@@ -79,19 +72,19 @@ CodeExecutionTool::CodeExecutionRequest::~CodeExecutionRequest() {
void CodeExecutionTool::CodeExecutionRequest::DidFinishLoad(
content::RenderFrameHost* render_frame_host,
const GURL& validated_url) {
if (!render_frame_host->GetParent() || wrapped_js_.empty()) {
if (!render_frame_host->GetParent() || script_.empty()) {
return;
}
render_frame_host->GetRemoteAssociatedInterfaces()->GetInterface(&injector_);
auto wrapped_js_utf16 = base::UTF8ToUTF16(wrapped_js_);
auto script_utf16 = base::UTF8ToUTF16(script_);
// Clear the wrapped script to avoid re-using it.
wrapped_js_ = {};
script_ = {};
injector_->RequestAsyncExecuteScript(
content::ISOLATED_WORLD_ID_GLOBAL, wrapped_js_utf16,
content::ISOLATED_WORLD_ID_GLOBAL, script_utf16,
blink::mojom::UserActivationOption::kActivate,
blink::mojom::PromiseResultOption::kAwait,
base::BindOnce(&CodeExecutionRequest::HandleResult,
@@ -109,29 +102,125 @@ void CodeExecutionTool::CodeExecutionRequest::OnDidAddMessageToConsole(
}
void CodeExecutionTool::CodeExecutionRequest::HandleResult(base::Value result) {
if (!result.is_bool() || !result.GetBool()) {
std::move(resolve_callback_).Run("Error: Syntax error");
if (!result.is_list()) {
std::move(resolve_callback_).Run("Error: Syntax error", {});
return;
}
std::move(resolve_callback_).Run(base::JoinString(console_logs_, "\n"));
std::string console_logs_str = base::JoinString(console_logs_, "\n");
std::move(resolve_callback_)
.Run(std::move(console_logs_str), std::move(result).TakeList());
}
void CodeExecutionTool::CodeExecutionRequest::HandleTimeout() {
std::move(resolve_callback_).Run("Error: Time limit exceeded");
std::move(resolve_callback_).Run("Error: Time limit exceeded", {});
}
void CodeExecutionTool::ResolveRequest(
std::list<CodeExecutionRequest>::iterator request_it,
UseToolCallback callback,
std::string output) {
std::string console_logs,
base::ListValue artifacts) {
requests_.erase(request_it);
std::move(callback).Run(CreateContentBlocksForText(output), {});
std::vector<mojom::ToolArtifactPtr> artifact_ptrs;
std::optional<std::string> error;
// Process artifacts
for (const auto& artifact : artifacts) {
const auto* artifact_dict = artifact.GetIfDict();
if (!artifact_dict) {
error = "Error: Artifact must be an object";
break;
}
const auto* type = artifact_dict->FindString(kArtifactTypeKey);
const auto* content = artifact_dict->Find(kArtifactContentKey);
if (!type || !content) {
error = "Error: Artifact missing required 'type' or 'content' field";
break;
}
// Find matching plugin and validate artifact
bool plugin_found = false;
for (const auto& plugin : code_plugins_) {
if (plugin->ArtifactType() != *type) {
continue;
}
plugin_found = true;
if (auto validation_error = plugin->ValidateArtifact(*content)) {
error = base::StrCat({"Error: ", *validation_error});
}
break;
}
if (!plugin_found) {
error =
base::StrCat({"Error: Artifact type '", *type, "' is not supported"});
break;
}
if (error) {
break;
}
// Serialize content to JSON string for storage
std::string content_json;
if (!base::JSONWriter::Write(*content, &content_json)) {
error = "Error: Failed to serialize artifact content";
break;
}
// Add artifact
artifact_ptrs.push_back(
mojom::ToolArtifact::New(*type, std::move(content_json)));
}
// Construct final content blocks
std::vector<mojom::ContentBlockPtr> content_blocks;
// If error occurred, use error message instead of console logs
if (error) {
content_blocks = CreateContentBlocksForText(std::move(*error));
artifact_ptrs.clear();
} else {
content_blocks = CreateContentBlocksForText(std::move(console_logs));
}
std::move(callback).Run(std::move(content_blocks), std::move(artifact_ptrs));
}
CodeExecutionTool::CodeExecutionTool(content::BrowserContext* browser_context)
: profile_(Profile::FromBrowserContext(browser_context)),
execution_time_limit_(kExecutionTimeLimit) {}
execution_time_limit_(kExecutionTimeLimit) {
// Build the description with plugin information
std::vector<std::string_view> plugin_descriptions;
for (const auto& plugin : code_plugins_) {
plugin_descriptions.emplace_back(plugin->Description());
}
tool_description_ = base::StrCat(
{"Execute JavaScript code and capture console output. "
"Use only when the task requires code execution for providing an "
"accurate answer. "
"Do not use this if you are able to answer without executing code. "
"Do not use this for content generation. "
"Do not use this for fetching information from the internet. "
"Use console.log() to output results. "
"The code will be executed in a sandboxed environment. "
"Network requests are not allowed. "
"bignumber.js is available in the global scope. Use it for any "
"decimal math (i.e. financial calculations). "
"Do not use require to import bignumber.js, as it is not needed. ",
base::JoinString(plugin_descriptions, " "),
"\nExample tasks that require code execution:\n"
" - Financial calculations (e.g. compound interest)\n"
" - Analyzing data or web content\n"
"Example tasks that do not require code execution:\n"
" - Very simple calculations (e.g. 2 + 2)\n"
" - Finding the 4th prime number\n"
" - Retrieving weather information for a location"});
}
CodeExecutionTool::~CodeExecutionTool() = default;
@@ -140,25 +229,7 @@ std::string_view CodeExecutionTool::Name() const {
}
std::string_view CodeExecutionTool::Description() const {
return "Execute JavaScript code and capture console output. "
"Use only when the task requires code execution for providing an "
"accurate answer. "
"Do not use this if you are able to answer without executing code. "
"Do not use this for content generation. "
"Do not use this for fetching information from the internet. "
"Use console.log() to output results. "
"The code will be executed in a sandboxed environment. "
"Network requests are not allowed. "
"bignumber.js is available in the global scope. Use it for any "
"decimal math (i.e. financial calculations). "
"Do not use require to import bignumber.js, as it is not needed.\n"
"Example tasks that require code execution:\n"
" - Financial calculations (e.g. compound interest)\n"
" - Analyzing data or web content\n"
"Example tasks that do not require code execution:\n"
" - Very simple calculations (e.g. 2 + 2)\n"
" - Finding the 4th prime number\n"
" - Retrieving weather information for a location";
return tool_description_;
}
std::optional<base::DictValue> CodeExecutionTool::InputProperties() const {
@@ -190,6 +261,30 @@ void CodeExecutionTool::SetExecutionTimeLimitForTesting(
execution_time_limit_ = time_limit;
}
void CodeExecutionTool::AddCodePluginForTesting(
std::unique_ptr<CodePlugin> plugin) {
code_plugins_.push_back(std::move(plugin));
}
std::string CodeExecutionTool::WrapScript(const std::string& script) const {
auto bignumber_js =
ui::ResourceBundle::GetSharedInstance().LoadDataResourceString(
IDR_AI_CHAT_BIGNUMBER_JS);
std::vector<std::string_view> plugin_scripts;
for (const auto& plugin : code_plugins_) {
if (script.find(plugin->InclusionKeyword()) != std::string::npos) {
plugin_scripts.push_back(plugin->SetupScript());
}
}
return base::StrCat({"(async function() { let codeExecArtifacts = []; ",
bignumber_js, base::StrCat(plugin_scripts), " try { ",
script,
" } catch (error) { console.error(error.toString()); } "
"return codeExecArtifacts; })()"});
}
void CodeExecutionTool::UseTool(const std::string& input_json,
UseToolCallback callback) {
auto input_dict = base::JSONReader::ReadDict(
@@ -211,7 +306,9 @@ void CodeExecutionTool::UseTool(const std::string& input_json,
return;
}
requests_.emplace_back(profile_, *script, execution_time_limit_);
std::string wrapped_script = WrapScript(*script);
requests_.emplace_back(profile_, std::move(wrapped_script),
execution_time_limit_);
auto request_it = std::prev(requests_.end());
request_it->SetResolveCallback(
browser/ai_chat/code_execution_tool.h → browser/ai_chat/tools/code_execution_tool.h
+15 -9
View File
@@ -3,8 +3,8 @@
// 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_CODE_EXECUTION_TOOL_H_
#define BRAVE_BROWSER_AI_CHAT_CODE_EXECUTION_TOOL_H_
#ifndef BRAVE_BROWSER_AI_CHAT_TOOLS_CODE_EXECUTION_TOOL_H_
#define BRAVE_BROWSER_AI_CHAT_TOOLS_CODE_EXECUTION_TOOL_H_
#include <list>
#include <memory>
@@ -17,6 +17,7 @@
#include "base/memory/weak_ptr.h"
#include "base/timer/timer.h"
#include "base/values.h"
#include "brave/components/ai_chat/core/browser/tools/code_plugin.h"
#include "brave/components/ai_chat/core/browser/tools/tool.h"
#include "brave/components/script_injector/common/mojom/script_injector.mojom.h"
#include "content/public/browser/web_contents_observer.h"
@@ -33,8 +34,6 @@ class RenderFrameHost;
namespace ai_chat {
class CodeSandboxWebContentsObserver;
// Tool for executing JavaScript code and returning console.log output.
// This tool is provided by the browser and allows AI assistants to run
// JavaScript code in a sandboxed environment.
@@ -62,14 +61,16 @@ class CodeExecutionTool : public Tool {
UseToolCallback callback) override;
void SetExecutionTimeLimitForTesting(base::TimeDelta time_limit);
void AddCodePluginForTesting(std::unique_ptr<CodePlugin> plugin);
private:
class CodeExecutionRequest : public content::WebContentsObserver {
public:
using ResolveCallback = base::OnceCallback<void(std::string)>;
using ResolveCallback = base::OnceCallback<void(std::string console_logs,
base::ListValue output)>;
CodeExecutionRequest(Profile* profile,
const std::string& script,
std::string script,
base::TimeDelta execution_time_limit);
~CodeExecutionRequest() override;
@@ -93,7 +94,7 @@ class CodeExecutionTool : public Tool {
void HandleTimeout();
std::unique_ptr<content::WebContents> web_contents_;
std::string wrapped_js_;
std::string script_;
mojo::AssociatedRemote<script_injector::mojom::ScriptInjector> injector_;
base::OneShotTimer timeout_timer_;
ResolveCallback resolve_callback_;
@@ -101,15 +102,20 @@ class CodeExecutionTool : public Tool {
base::WeakPtrFactory<CodeExecutionRequest> weak_ptr_factory_{this};
};
std::string WrapScript(const std::string& script) const;
void ResolveRequest(std::list<CodeExecutionRequest>::iterator request_it,
UseToolCallback callback,
std::string output);
std::string console_logs,
base::ListValue output);
raw_ptr<Profile> profile_;
std::vector<std::unique_ptr<CodePlugin>> code_plugins_;
std::string tool_description_;
std::list<CodeExecutionRequest> requests_;
base::TimeDelta execution_time_limit_;
};
} // namespace ai_chat
#endif // BRAVE_BROWSER_AI_CHAT_CODE_EXECUTION_TOOL_H_
#endif // BRAVE_BROWSER_AI_CHAT_TOOLS_CODE_EXECUTION_TOOL_H_
components/ai_chat/core/browser/BUILD.gn
+2
View File
@@ -73,6 +73,8 @@ static_library("browser") {
"skills_metrics.h",
"tab_tracker_service.cc",
"tab_tracker_service.h",
"tools/code_plugin.cc",
"tools/code_plugin.h",
"tools/memory_storage_tool.cc",
"tools/memory_storage_tool.h",
"tools/tool.cc",
components/ai_chat/core/browser/tools/code_plugin.cc
+17
View File
@@ -0,0 +1,17 @@
// 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/components/ai_chat/core/browser/tools/code_plugin.h"
namespace ai_chat {
CodePlugin::~CodePlugin() = default;
std::optional<std::string> CodePlugin::ValidateArtifact(
const base::Value& artifact_value) const {
return std::nullopt;
}
} // namespace ai_chat
components/ai_chat/core/browser/tools/code_plugin.h
+44
View File
@@ -0,0 +1,44 @@
// 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_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_CODE_PLUGIN_H_
#define BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_CODE_PLUGIN_H_
#include <optional>
#include <string>
#include <string_view>
#include "base/values.h"
namespace ai_chat {
// Interface for code execution plugins that provide additional utilities
// to the JavaScript execution environment.
class CodePlugin {
public:
virtual ~CodePlugin();
// Description of the plugin's capabilities for the tool description
virtual std::string_view Description() const = 0;
// Keyword that triggers inclusion of this plugin's setup script
virtual std::string_view InclusionKeyword() const = 0;
// JavaScript setup script to inject into the execution environment
virtual std::string_view SetupScript() const = 0;
// The artifact type this plugin handles
virtual std::string_view ArtifactType() const = 0;
// Validates an artifact from script execution. Returns an error message
// if validation fails, or std::nullopt if validation succeeds.
// |artifact_value| is the parsed JSON.
virtual std::optional<std::string> ValidateArtifact(
const base::Value& artifact_value) const;
};
} // namespace ai_chat
#endif // BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_CODE_PLUGIN_H_