diff --git a/browser/ai_chat/BUILD.gn b/browser/ai_chat/BUILD.gn index 62e159e5fc6..bd69c7856cf 100644 --- a/browser/ai_chat/BUILD.gn +++ b/browser/ai_chat/BUILD.gn @@ -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", ] diff --git a/browser/ai_chat/browser_tool_provider.cc b/browser/ai_chat/browser_tool_provider.cc index 0085d51ea6e..1df773f2ce7 100644 --- a/browser/ai_chat/browser_tool_provider.cc +++ b/browser/ai_chat/browser_tool_provider.cc @@ -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" diff --git a/browser/ai_chat/code_execution_tool_browsertest.cc b/browser/ai_chat/code_execution_tool_browsertest.cc index 01fca15d8f9..8ce35f04140 100644 --- a/browser/ai_chat/code_execution_tool_browsertest.cc +++ b/browser/ai_chat/code_execution_tool_browsertest.cc @@ -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 +#include #include +#include #include #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 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 ValidateArtifact( + const base::Value& artifact_value) const override { + return validation_error_; + } + + private: + std::optional 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* artifacts = nullptr) { base::RunLoop run_loop; tool_->UseTool( input_json, base::BindLambdaForTesting( - [&run_loop, output](std::vector result, - std::vector) { + [&run_loop, output, artifacts]( + std::vector result, + std::vector 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* 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()); + + std::string output; + std::vector 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(); + plugin->SetValidationError("content must be a number"); + tool_->AddCodePluginForTesting(std::move(plugin)); + + std::string output; + std::vector 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 diff --git a/browser/ai_chat/code_execution_tool.cc b/browser/ai_chat/tools/code_execution_tool.cc similarity index 54% rename from browser/ai_chat/code_execution_tool.cc rename to browser/ai_chat/tools/code_execution_tool.cc index c637166ebdb..ed970a31104 100644 --- a/browser/ai_chat/code_execution_tool.cc +++ b/browser/ai_chat/tools/code_execution_tool.cc @@ -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 #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::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 artifact_ptrs; + std::optional 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 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 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 CodeExecutionTool::InputProperties() const { @@ -190,6 +261,30 @@ void CodeExecutionTool::SetExecutionTimeLimitForTesting( execution_time_limit_ = time_limit; } +void CodeExecutionTool::AddCodePluginForTesting( + std::unique_ptr 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 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( diff --git a/browser/ai_chat/code_execution_tool.h b/browser/ai_chat/tools/code_execution_tool.h similarity index 82% rename from browser/ai_chat/code_execution_tool.h rename to browser/ai_chat/tools/code_execution_tool.h index 0e1e3b5494b..87cd6c4c6be 100644 --- a/browser/ai_chat/code_execution_tool.h +++ b/browser/ai_chat/tools/code_execution_tool.h @@ -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 #include @@ -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 plugin); private: class CodeExecutionRequest : public content::WebContentsObserver { public: - using ResolveCallback = base::OnceCallback; + using ResolveCallback = base::OnceCallback; 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 web_contents_; - std::string wrapped_js_; + std::string script_; mojo::AssociatedRemote injector_; base::OneShotTimer timeout_timer_; ResolveCallback resolve_callback_; @@ -101,15 +102,20 @@ class CodeExecutionTool : public Tool { base::WeakPtrFactory weak_ptr_factory_{this}; }; + std::string WrapScript(const std::string& script) const; + void ResolveRequest(std::list::iterator request_it, UseToolCallback callback, - std::string output); + std::string console_logs, + base::ListValue output); raw_ptr profile_; + std::vector> code_plugins_; + std::string tool_description_; std::list 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_ diff --git a/components/ai_chat/core/browser/BUILD.gn b/components/ai_chat/core/browser/BUILD.gn index a446aa7209f..e31d62bfb8f 100644 --- a/components/ai_chat/core/browser/BUILD.gn +++ b/components/ai_chat/core/browser/BUILD.gn @@ -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", diff --git a/components/ai_chat/core/browser/tools/code_plugin.cc b/components/ai_chat/core/browser/tools/code_plugin.cc new file mode 100644 index 00000000000..7e324a4020e --- /dev/null +++ b/components/ai_chat/core/browser/tools/code_plugin.cc @@ -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 CodePlugin::ValidateArtifact( + const base::Value& artifact_value) const { + return std::nullopt; +} + +} // namespace ai_chat diff --git a/components/ai_chat/core/browser/tools/code_plugin.h b/components/ai_chat/core/browser/tools/code_plugin.h new file mode 100644 index 00000000000..23bfa1bca92 --- /dev/null +++ b/components/ai_chat/core/browser/tools/code_plugin.h @@ -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 +#include +#include + +#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 ValidateArtifact( + const base::Value& artifact_value) const; +}; + +} // namespace ai_chat + +#endif // BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_CODE_PLUGIN_H_