Reland "[AI Chat] Support text file uploads with renderer-based extraction"" (#35613)

Revert "Revert "[AI Chat] Support text file uploads with renderer-based extra…"

This reverts commit c0e4f7effd.
This commit is contained in:
Anthony Tseng
2026-04-17 12:06:01 -07:00
committed by GitHub
parent 0a7437712c
commit 9df0bef8ab
38 changed files with 1268 additions and 219 deletions
+15 -1
View File
@@ -66,6 +66,7 @@ static_library("ai_chat") {
"//chrome/browser/actor",
"//chrome/browser/profiles:profile",
"//chrome/browser/ui/tabs:tabs_public",
"//chrome/common",
"//chrome/common:channel_info",
"//components/browsing_data/core",
"//components/keyed_service/content",
@@ -76,6 +77,7 @@ static_library("ai_chat") {
"//printing:printing_base",
"//services/data_decoder/public/cpp",
"//services/network/public/cpp",
"//third_party/blink/public/common",
"//ui/base",
"//ui/shell_dialogs",
]
@@ -84,7 +86,14 @@ static_library("ai_chat") {
deps += [ "//brave/build/android:jni_headers" ]
}
if (!is_android) {
deps += [ "//chrome/browser/ui/side_panel" ]
sources += [
"text_file_extractor.cc",
"text_file_extractor.h",
]
deps += [
"//chrome/browser/ui/side_panel",
"//third_party/blink/public/strings",
]
}
if (enable_pdf) {
@@ -289,6 +298,10 @@ source_set("unit_tests") {
sources += [ "tools/tab_management_tool_unittest.cc" ]
}
if (!is_android) {
sources += [ "text_file_extractor_unittest.cc" ]
}
if (enable_pdf) {
sources += [ "pdf_text_extractor_unittest.cc" ]
}
@@ -357,6 +370,7 @@ source_set("browser_tests") {
"ai_chat_ui_browsertest.cc",
"code_execution_tool_browsertest.cc",
"page_content_fetcher_browsertest.cc",
"text_file_extractor_browsertest.cc",
]
if (enable_pdf) {
+1
View File
@@ -2,6 +2,7 @@ include_rules = [
"+brave/components/restricted_web_contents_delegate",
"+brave/services/printing/public/mojom",
"+brave/components/text_recognition/common",
"+third_party/blink/public/strings",
]
specific_include_rules = {
+25 -2
View File
@@ -64,11 +64,34 @@ FileTextExtractorBase::~FileTextExtractorBase() {
Cleanup();
}
void FileTextExtractorBase::ExtractText(
content::BrowserContext* browser_context,
const base::FilePath& file_path,
ExtractTextCallback callback) {
CHECK(!callback_) << "ExtractText called while extraction in progress";
callback_ = std::move(callback);
LoadInWebContents(browser_context, file_path);
}
void FileTextExtractorBase::ExtractText(
content::BrowserContext* browser_context,
std::vector<uint8_t> file_bytes,
const base::FilePath::StringType& extension,
ExtractTextCallback callback) {
CHECK(!callback_) << "ExtractText called while extraction in progress";
callback_ = std::move(callback);
WriteTempFileAndLoad(browser_context, std::move(file_bytes), extension);
}
network::mojom::WebSandboxFlags
FileTextExtractorBase::AdditionalUnsandboxFlags() const {
return network::mojom::WebSandboxFlags::kNone;
}
GURL FileTextExtractorBase::GetLoadURL(const base::FilePath& file_path) const {
return net::FilePathToFileURL(file_path);
}
void FileTextExtractorBase::LoadInWebContents(
content::BrowserContext* browser_context,
const base::FilePath& file_path) {
@@ -89,8 +112,8 @@ void FileTextExtractorBase::LoadInWebContents(
base::BindOnce(&FileTextExtractorBase::OnTimeout,
base::Unretained(this)));
const GURL file_url = net::FilePathToFileURL(file_path);
web_contents_->GetController().LoadURL(file_url, content::Referrer(),
const GURL load_url = GetLoadURL(file_path);
web_contents_->GetController().LoadURL(load_url, content::Referrer(),
ui::PAGE_TRANSITION_AUTO_TOPLEVEL,
std::string());
}
@@ -18,6 +18,7 @@
#include "brave/components/restricted_web_contents_delegate/restricted_web_contents_delegate.h"
#include "content/public/browser/web_contents_observer.h"
#include "services/network/public/cpp/web_sandbox_flags.h"
#include "url/gurl.h"
namespace content {
class BrowserContext;
@@ -44,6 +45,21 @@ class FileTextExtractorBase : public RestrictedWebContentsDelegate,
FileTextExtractorBase(const FileTextExtractorBase&) = delete;
FileTextExtractorBase& operator=(const FileTextExtractorBase&) = delete;
// Two entry points for text extraction:
// Use an existing file path directly (e.g. from file picker).
void ExtractText(content::BrowserContext* browser_context,
const base::FilePath& file_path,
ExtractTextCallback callback);
// Write bytes to a temp file first (e.g. from drag-and-drop).
// |extension| is the file extension for MIME type detection (without
// leading dot).
void ExtractText(content::BrowserContext* browser_context,
std::vector<uint8_t> file_bytes,
const base::FilePath::StringType& extension,
ExtractTextCallback callback);
protected:
// Called when the document has loaded and is ready for text extraction.
// Subclasses must implement this and call Finish() with the result.
@@ -53,6 +69,10 @@ class FileTextExtractorBase : public RestrictedWebContentsDelegate,
// (Scripts, Origin, Navigation). Override to add more, e.g. kPlugins.
virtual network::mojom::WebSandboxFlags AdditionalUnsandboxFlags() const;
// Returns the URL to load for the given file path. Default returns a
// file:// URL. Override to customize, e.g. view-source:file://.
virtual GURL GetLoadURL(const base::FilePath& file_path) const;
// Starts loading a file in a hidden WebContents.
void LoadInWebContents(content::BrowserContext* browser_context,
const base::FilePath& file_path);
-17
View File
@@ -19,23 +19,6 @@ PdfTextExtractor::PdfTextExtractor() = default;
PdfTextExtractor::~PdfTextExtractor() = default;
void PdfTextExtractor::ExtractText(content::BrowserContext* browser_context,
const base::FilePath& pdf_path,
ExtractTextCallback callback) {
CHECK(!callback_) << "ExtractText called while extraction in progress";
callback_ = std::move(callback);
LoadInWebContents(browser_context, pdf_path);
}
void PdfTextExtractor::ExtractText(content::BrowserContext* browser_context,
std::vector<uint8_t> pdf_bytes,
ExtractTextCallback callback) {
CHECK(!callback_) << "ExtractText called while extraction in progress";
callback_ = std::move(callback);
WriteTempFileAndLoad(browser_context, std::move(pdf_bytes),
FILE_PATH_LITERAL("pdf"));
}
network::mojom::WebSandboxFlags PdfTextExtractor::AdditionalUnsandboxFlags()
const {
// Plugins are required for the PDF viewer MimeHandlerView.
-21
View File
@@ -6,16 +6,11 @@
#ifndef BRAVE_BROWSER_AI_CHAT_PDF_TEXT_EXTRACTOR_H_
#define BRAVE_BROWSER_AI_CHAT_PDF_TEXT_EXTRACTOR_H_
#include <optional>
#include <string>
#include <vector>
#include "base/memory/weak_ptr.h"
#include "brave/browser/ai_chat/file_text_extractor_base.h"
#include "services/network/public/cpp/web_sandbox_flags.h"
namespace content {
class BrowserContext;
class RenderFrameHost;
} // namespace content
@@ -25,28 +20,12 @@ namespace ai_chat {
// The PDF viewer extension + ScreenAI OCR pipeline runs, then page text is
// extracted via PDFDocumentHelper::GetPageText().
//
// Two entry points:
// - ExtractText(browser_context, pdf_path, callback)
// Uses an existing file path directly (e.g. from file picker).
// - ExtractText(browser_context, pdf_bytes, callback)
// Writes bytes to a temp file first (e.g. from drag-and-drop).
//
// The extractor should be kept alive until the callback fires.
class PdfTextExtractor : public FileTextExtractorBase {
public:
PdfTextExtractor();
~PdfTextExtractor() override;
// Use an existing file path directly (no temp file created).
void ExtractText(content::BrowserContext* browser_context,
const base::FilePath& pdf_path,
ExtractTextCallback callback);
// Write bytes to a temp file first, then extract.
void ExtractText(content::BrowserContext* browser_context,
std::vector<uint8_t> pdf_bytes,
ExtractTextCallback callback);
private:
// FileTextExtractorBase:
void OnDocumentReady() override;
@@ -60,7 +60,7 @@ IN_PROC_BROWSER_TEST_F(PdfTextExtractorBrowserTest,
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), std::move(*pdf_bytes),
future.GetCallback());
FILE_PATH_LITERAL("pdf"), future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result.has_value());
@@ -40,7 +40,7 @@ TEST_F(PdfTextExtractorTest, BytesOverload_TimeoutReturnsNullopt) {
std::vector<uint8_t> dummy_pdf = {0x25, 0x50, 0x44, 0x46}; // %PDF
extractor->ExtractText(browser_context(), std::move(dummy_pdf),
future.GetCallback());
FILE_PATH_LITERAL("pdf"), future.GetCallback());
// Fast-forward past the 30s extraction timeout.
// This also processes pending ThreadPool tasks (temp-file write).
@@ -99,7 +99,7 @@ TEST_F(PdfTextExtractorTest, BytesOverload_CleanupAfterTimeout) {
std::vector<uint8_t> dummy_pdf = {0x25, 0x50, 0x44, 0x46};
extractor->ExtractText(browser_context(), std::move(dummy_pdf),
future.GetCallback());
FILE_PATH_LITERAL("pdf"), future.GetCallback());
// Fast-forward past timeout to trigger cleanup.
// This also processes pending ThreadPool tasks (temp-file write).
+71
View File
@@ -0,0 +1,71 @@
// 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/text_file_extractor.h"
#include <utility>
#include "base/functional/bind.h"
#include "base/logging.h"
#include "base/strings/strcat.h"
#include "base/strings/string_util.h"
#include "chrome/common/chrome_isolated_world_ids.h"
#include "content/public/browser/render_frame_host.h"
#include "content/public/browser/web_contents.h"
#include "content/public/common/url_constants.h"
#include "net/base/filename_util.h"
#include "third_party/blink/public/strings/grit/blink_strings.h"
#include "ui/base/l10n/l10n_util.h"
namespace ai_chat {
TextFileExtractor::TextFileExtractor() = default;
TextFileExtractor::~TextFileExtractor() = default;
GURL TextFileExtractor::GetLoadURL(const base::FilePath& file_path) const {
// Use view-source: to prevent HTML/XHTML from being rendered (which could
// execute scripts or load external resources). The source is displayed as
// raw text in the view-source document.
return GURL(base::StrCat({content::kViewSourceScheme, ":",
net::FilePathToFileURL(file_path).spec()}));
}
void TextFileExtractor::OnDocumentReady() {
auto* rfh = GetWebContents()->GetPrimaryMainFrame();
if (!rfh) {
Finish(std::nullopt);
return;
}
rfh->ExecuteJavaScriptInIsolatedWorld(
u"document.body.innerText",
base::BindOnce(&TextFileExtractor::OnTextExtracted,
weak_ptr_factory_.GetWeakPtr()),
ISOLATED_WORLD_ID_BRAVE_INTERNAL);
}
void TextFileExtractor::OnTextExtracted(base::Value result) {
if (result.is_string()) {
// view-source: innerText starts with the localized "Line wrap" label
// followed by a newline, and indents each line with a leading tab.
// Strip the label prefix and the per-line tabs.
// For non-empty files the prefix is "Line wrap\n\t", for empty files
// it's just "Line wrap" with no trailing newline or tab.
std::string text = std::move(result).TakeString();
std::string line_wrap_label =
l10n_util::GetStringUTF8(IDS_VIEW_SOURCE_LINE_WRAP);
if (auto rest = base::RemovePrefix(text, line_wrap_label)) {
text = std::string(base::RemovePrefix(*rest, "\n\t").value_or(*rest));
}
base::ReplaceSubstringsAfterOffset(&text, 0, "\n\t", "\n");
Finish(std::move(text));
} else {
DVLOG(1) << "TextFileExtractor: JS returned non-string result";
Finish(std::nullopt);
}
}
} // namespace ai_chat
+45
View File
@@ -0,0 +1,45 @@
// 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_TEXT_FILE_EXTRACTOR_H_
#define BRAVE_BROWSER_AI_CHAT_TEXT_FILE_EXTRACTOR_H_
#include "base/gtest_prod_util.h"
#include "base/memory/weak_ptr.h"
#include "brave/browser/ai_chat/file_text_extractor_base.h"
namespace ai_chat {
// Extracts text from a file by loading it in a hidden background WebContents.
// Chromium's renderer handles MIME sniffing and renders the file content,
// then text is extracted via document.body.innerText.
//
// The extractor should be kept alive until the callback fires.
class TextFileExtractor : public FileTextExtractorBase {
public:
TextFileExtractor();
~TextFileExtractor() override;
private:
FRIEND_TEST_ALL_PREFIXES(TextFileExtractorTest,
OnTextExtracted_StripsViewSourcePrefix);
FRIEND_TEST_ALL_PREFIXES(TextFileExtractorTest,
OnTextExtracted_StripsLeadingTabs);
FRIEND_TEST_ALL_PREFIXES(TextFileExtractorTest, OnTextExtracted_EmptyFile);
FRIEND_TEST_ALL_PREFIXES(TextFileExtractorTest,
OnTextExtracted_NonStringReturnsNullopt);
// FileTextExtractorBase:
GURL GetLoadURL(const base::FilePath& file_path) const override;
void OnDocumentReady() override;
void OnTextExtracted(base::Value result);
base::WeakPtrFactory<TextFileExtractor> weak_ptr_factory_{this};
};
} // namespace ai_chat
#endif // BRAVE_BROWSER_AI_CHAT_TEXT_FILE_EXTRACTOR_H_
@@ -0,0 +1,139 @@
// 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/text_file_extractor.h"
#include <optional>
#include <string>
#include <vector>
#include "base/files/file_util.h"
#include "base/path_service.h"
#include "base/test/test_future.h"
#include "base/threading/thread_restrictions.h"
#include "brave/components/constants/brave_paths.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/browser/ui/browser.h"
#include "chrome/test/base/in_process_browser_test.h"
#include "content/public/test/browser_test.h"
#include "testing/gtest/include/gtest/gtest.h"
namespace ai_chat {
// Expected content of test/data/leo/dummy.txt
constexpr char kExpectedTextContent[] =
"Hello from a text file.\nThis is line two.";
class TextFileExtractorBrowserTest : public InProcessBrowserTest {
protected:
base::FilePath GetTestFilePath(std::string_view filename) {
base::ScopedAllowBlockingForTesting allow_blocking;
return base::PathService::CheckedGet(brave::DIR_TEST_DATA)
.AppendASCII("leo")
.AppendASCII(filename);
}
content::BrowserContext* browser_context() { return browser()->profile(); }
};
IN_PROC_BROWSER_TEST_F(TextFileExtractorBrowserTest,
PathOverload_ExtractsText) {
base::FilePath txt_path = GetTestFilePath("dummy.txt");
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), txt_path, future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_EQ(*result, kExpectedTextContent);
}
IN_PROC_BROWSER_TEST_F(TextFileExtractorBrowserTest,
BytesOverload_ExtractsText) {
std::optional<std::vector<uint8_t>> file_bytes;
{
base::ScopedAllowBlockingForTesting allow_blocking;
file_bytes = base::ReadFileToBytes(GetTestFilePath("dummy.txt"));
}
ASSERT_TRUE(file_bytes.has_value());
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), std::move(*file_bytes),
FILE_PATH_LITERAL("txt"), future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_EQ(*result, kExpectedTextContent);
}
// HTML files must not be rendered — the extracted text should be the raw
// source containing HTML tags, not the rendered text content.
IN_PROC_BROWSER_TEST_F(TextFileExtractorBrowserTest,
PathOverload_HtmlNotRendered) {
base::FilePath html_path = GetTestFilePath("dummy.html");
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), html_path, future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result.has_value());
// Raw source must contain HTML tags — proves it was NOT rendered.
EXPECT_NE(result->find("<p>Hello from an HTML file.</p>"), std::string::npos);
EXPECT_NE(result->find("<script>"), std::string::npos);
}
IN_PROC_BROWSER_TEST_F(TextFileExtractorBrowserTest,
BytesOverload_HtmlNotRendered) {
std::optional<std::vector<uint8_t>> file_bytes;
{
base::ScopedAllowBlockingForTesting allow_blocking;
file_bytes = base::ReadFileToBytes(GetTestFilePath("dummy.html"));
}
ASSERT_TRUE(file_bytes.has_value());
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), std::move(*file_bytes),
FILE_PATH_LITERAL("html"), future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_NE(result->find("<p>Hello from an HTML file.</p>"), std::string::npos);
EXPECT_NE(result->find("<script>"), std::string::npos);
}
IN_PROC_BROWSER_TEST_F(TextFileExtractorBrowserTest, EmptyFile_ReturnsEmpty) {
// Create a temporary empty file
base::FilePath temp_path;
{
base::ScopedAllowBlockingForTesting allow_blocking;
ASSERT_TRUE(base::CreateTemporaryFile(&temp_path));
}
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), temp_path, future.GetCallback());
auto result = future.Take();
// Empty file should either return nullopt or empty string.
if (result.has_value()) {
EXPECT_TRUE(result->empty());
}
{
base::ScopedAllowBlockingForTesting allow_blocking;
base::DeleteFile(temp_path);
}
}
} // namespace ai_chat
@@ -0,0 +1,184 @@
// 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/text_file_extractor.h"
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "base/files/file_util.h"
#include "base/files/scoped_temp_dir.h"
#include "base/strings/strcat.h"
#include "base/test/test_future.h"
#include "chrome/test/base/testing_profile.h"
#include "content/public/test/test_renderer_host.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/blink/public/strings/grit/blink_strings.h"
#include "ui/base/l10n/l10n_util.h"
namespace ai_chat {
class TextFileExtractorTest : public content::RenderViewHostTestHarness {
public:
TextFileExtractorTest()
: content::RenderViewHostTestHarness(
base::test::TaskEnvironment::TimeSource::MOCK_TIME) {}
std::unique_ptr<content::BrowserContext> CreateBrowserContext() override {
return std::make_unique<TestingProfile>();
}
};
// Without a real renderer producing text content, the extraction should
// time out and return nullopt.
TEST_F(TextFileExtractorTest, BytesOverload_TimeoutReturnsNullopt) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
std::vector<uint8_t> text_bytes = {'h', 'e', 'l', 'l', 'o'};
extractor->ExtractText(browser_context(), std::move(text_bytes),
FILE_PATH_LITERAL("txt"), future.GetCallback());
task_environment()->FastForwardBy(base::Seconds(31));
auto result = future.Take();
EXPECT_FALSE(result.has_value());
}
// Same timeout test but using the file-path overload (no temp file).
TEST_F(TextFileExtractorTest, PathOverload_TimeoutReturnsNullopt) {
base::ScopedTempDir temp_dir;
ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
base::FilePath txt_path = temp_dir.GetPath().AppendASCII("test.txt");
ASSERT_TRUE(base::WriteFile(txt_path, "hello world"));
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), txt_path, future.GetCallback());
task_environment()->FastForwardBy(base::Seconds(31));
auto result = future.Take();
EXPECT_FALSE(result.has_value());
}
// Destroying the extractor while an extraction is in-flight must not crash
// or leak.
TEST_F(TextFileExtractorTest, DestroyDuringExtraction_NoCrash) {
auto extractor = std::make_unique<TextFileExtractor>();
base::ScopedTempDir temp_dir;
ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
base::FilePath txt_path = temp_dir.GetPath().AppendASCII("test.txt");
ASSERT_TRUE(base::WriteFile(txt_path, "hello world"));
bool callback_called = false;
extractor->ExtractText(
browser_context(), txt_path,
base::BindOnce(
[](bool* called, std::optional<std::string>) { *called = true; },
&callback_called));
// Destroy while extraction is in progress — should not crash.
extractor.reset();
EXPECT_FALSE(callback_called);
}
// The bytes overload writes to a temp file. Verify cleanup completes
// without crashing after timeout.
TEST_F(TextFileExtractorTest, BytesOverload_CleanupAfterTimeout) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
std::vector<uint8_t> text_bytes = {'h', 'e', 'l', 'l', 'o'};
extractor->ExtractText(browser_context(), std::move(text_bytes),
FILE_PATH_LITERAL("txt"), future.GetCallback());
task_environment()->FastForwardBy(base::Seconds(31));
ASSERT_TRUE(future.Wait());
}
// The file-path overload should NOT delete the original file after extraction.
TEST_F(TextFileExtractorTest, PathOverload_OriginalFileNotDeleted) {
base::ScopedTempDir temp_dir;
ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
base::FilePath txt_path = temp_dir.GetPath().AppendASCII("keep_me.txt");
ASSERT_TRUE(base::WriteFile(txt_path, "this file should persist"));
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->ExtractText(browser_context(), txt_path, future.GetCallback());
task_environment()->FastForwardBy(base::Seconds(31));
ASSERT_TRUE(future.Wait());
// The original file must still exist — only temp files are cleaned up.
EXPECT_TRUE(base::PathExists(txt_path));
}
// Unit tests for OnTextExtracted view-source stripping logic.
TEST_F(TextFileExtractorTest, OnTextExtracted_StripsViewSourcePrefix) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->callback_ = future.GetCallback();
// Simulate view-source innerText: "<Line wrap label>\n\tline1\n\tline2"
std::string line_wrap = l10n_util::GetStringUTF8(IDS_VIEW_SOURCE_LINE_WRAP);
extractor->OnTextExtracted(
base::Value(base::StrCat({line_wrap, "\n\tHello world\n\tSecond line"})));
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_EQ(*result, "Hello world\nSecond line");
}
TEST_F(TextFileExtractorTest, OnTextExtracted_StripsLeadingTabs) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->callback_ = future.GetCallback();
// Tabs embedded within content (not at line start) should be preserved.
std::string line_wrap = l10n_util::GetStringUTF8(IDS_VIEW_SOURCE_LINE_WRAP);
extractor->OnTextExtracted(
base::Value(base::StrCat({line_wrap, "\n\tcol1\tcol2\n\tcol3\tcol4"})));
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_EQ(*result, "col1\tcol2\ncol3\tcol4");
}
TEST_F(TextFileExtractorTest, OnTextExtracted_EmptyFile) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->callback_ = future.GetCallback();
// Empty file in view-source returns just the label with no newline.
std::string line_wrap = l10n_util::GetStringUTF8(IDS_VIEW_SOURCE_LINE_WRAP);
extractor->OnTextExtracted(base::Value(line_wrap));
auto result = future.Take();
ASSERT_TRUE(result.has_value());
EXPECT_TRUE(result->empty());
}
TEST_F(TextFileExtractorTest, OnTextExtracted_NonStringReturnsNullopt) {
auto extractor = std::make_unique<TextFileExtractor>();
base::test::TestFuture<std::optional<std::string>> future;
extractor->callback_ = future.GetCallback();
extractor->OnTextExtracted(base::Value(42));
auto result = future.Take();
EXPECT_FALSE(result.has_value());
}
} // namespace ai_chat
+52 -19
View File
@@ -12,16 +12,18 @@
#include "base/files/file.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/no_destructor.h"
#include "base/task/thread_pool.h"
#include "brave/components/ai_chat/core/browser/utils.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/common.mojom.h"
#include "chrome/browser/profiles/profile.h"
#include "content/public/browser/web_contents.h"
#include "media/base/mime_util.h"
#include "net/base/mime_util.h"
#include "printing/printing_utils.h"
#include "services/data_decoder/public/cpp/data_decoder.h"
#include "services/data_decoder/public/cpp/decode_image.h"
#include "third_party/blink/public/common/mime_util/mime_util.h"
#include "third_party/skia/include/core/SkBitmap.h"
#include "ui/gfx/codec/png_codec.h"
#include "ui/shell_dialogs/selected_file_info.h"
@@ -51,8 +53,29 @@ std::optional<mojom::UploadedFileType> DetermineFileType(
return mojom::UploadedFileType::kImage;
}
// If no recognized extension, return nullopt
return std::nullopt;
// If the extension has a known MIME type, reject clearly binary types
// (e.g. zip, exe, tar) while accepting text-renderable types.
auto extension = file_path.Extension();
if (!extension.empty()) {
std::string mime_type;
if (net::GetMimeTypeFromExtension(extension.substr(1), &mime_type)) {
// Reject known binary MIME types (media, archives, executables, etc.)
if (blink::IsSupportedImageMimeType(mime_type) ||
media::IsSupportedMediaMimeType(mime_type)) {
return std::nullopt;
}
// Accept types that Chromium can render as text
if (blink::IsSupportedNonImageMimeType(mime_type)) {
return mojom::UploadedFileType::kText;
}
// Known MIME type but not renderable (e.g. application/zip)
return std::nullopt;
}
}
// No MIME mapping for this extension (or no extension). Let the renderer
// try via MIME sniffing, matching how Chromium handles file:// URLs.
return mojom::UploadedFileType::kText;
}
// base::ReadFileToBytes doesn't handle content uri so we need to read from
@@ -129,9 +152,6 @@ void UploadFileHelper::UploadFile(std::unique_ptr<ui::SelectFilePolicy> policy,
select_file_dialog_ = ui::SelectFileDialog::Create(this, std::move(policy));
ui::SelectFileDialog::FileTypeInfo info;
info.allowed_paths = ui::SelectFileDialog::FileTypeInfo::NATIVE_PATH;
info.extensions = {{FILE_PATH_LITERAL("png"), FILE_PATH_LITERAL("jpeg"),
FILE_PATH_LITERAL("jpg"), FILE_PATH_LITERAL("webp"),
FILE_PATH_LITERAL("pdf")}};
#if BUILDFLAG(IS_ANDROID)
// Set the list of acceptable MIME types for the file picker; this will apply
// to any subsequent SelectFile() calls.
@@ -208,12 +228,16 @@ void UploadFileHelper::MultiFilesSelected(
uploaded_files.push_back(mojom::UploadedFile::New(
filepath.AsUTF8Unsafe(), file_data->size(), *file_data,
*file_type_opt, std::nullopt));
} else {
// Include rejected files as empty stubs so the frontend
// can detect dropped files and show an error alert.
uploaded_files.push_back(mojom::UploadedFile::New(
filepath.AsUTF8Unsafe(), 0, std::vector<uint8_t>(),
mojom::UploadedFileType::kText, std::nullopt));
}
}
std::move(callback).Run(
uploaded_files.empty()
? std::nullopt
: std::make_optional(std::move(uploaded_files)));
std::make_optional(std::move(uploaded_files)));
},
std::move(upload_file_callback_)));
@@ -241,8 +265,10 @@ void UploadFileHelper::MultiFilesSelected(
auto file_type_opt = DetermineFileType(filepath, *file_data);
if (file_type_opt &&
*file_type_opt == mojom::UploadedFileType::kPdf) {
// For PDFs, just return the raw data without processing
(*file_type_opt == mojom::UploadedFileType::kPdf ||
*file_type_opt == mojom::UploadedFileType::kText)) {
// For PDFs and text files, return raw data without processing.
// Text extraction happens via ProcessPdfFile/ProcessTextFile.
std::move(callback).Run(std::make_tuple(
std::move(file_data), std::move(filepath), file_type_opt));
} else if (file_type_opt &&
@@ -297,13 +323,14 @@ void UploadFileHelper::OnFileRead(
// Determine file type based on extension and validate PDF content
auto file_type_opt = DetermineFileType(std::get<1>(result), *file_data);
if (file_type_opt && *file_type_opt == mojom::UploadedFileType::kPdf) {
// Return raw PDF data; text extraction happens via ProcessPdfFile mojo
// endpoint when uploading from WebUI drag-and-drop.
if (file_type_opt && (*file_type_opt == mojom::UploadedFileType::kPdf ||
*file_type_opt == mojom::UploadedFileType::kText)) {
// Return raw data; text extraction happens via ProcessPdfFile or
// ProcessTextFile mojo endpoint.
std::vector<mojom::UploadedFilePtr> files;
files.push_back(mojom::UploadedFile::New(
std::get<1>(result).AsUTF8Unsafe(), file_data->size(), *file_data,
mojom::UploadedFileType::kPdf, std::nullopt));
files.push_back(mojom::UploadedFile::New(std::get<1>(result).AsUTF8Unsafe(),
file_data->size(), *file_data,
*file_type_opt, std::nullopt));
std::move(upload_file_callback_).Run(std::make_optional(std::move(files)));
} else if (file_type_opt &&
*file_type_opt == mojom::UploadedFileType::kImage) {
@@ -314,8 +341,14 @@ void UploadFileHelper::OnFileRead(
weak_ptr_factory_.GetWeakPtr(),
std::get<1>(result).AsUTF8Unsafe()));
} else {
// Fail if we cannot handle this file type
std::move(upload_file_callback_).Run(std::nullopt);
// Include as empty stub so the frontend can detect the unsupported
// file and show an error (as opposed to nullopt which means the user
// cancelled the file picker).
std::vector<mojom::UploadedFilePtr> files;
files.push_back(mojom::UploadedFile::New(
std::get<1>(result).AsUTF8Unsafe(), 0, std::vector<uint8_t>(),
mojom::UploadedFileType::kText, std::nullopt));
std::move(upload_file_callback_).Run(std::make_optional(std::move(files)));
}
}
+159 -15
View File
@@ -135,18 +135,10 @@ TEST_F(UploadFileHelperTest, AcceptedFileExtensions) {
EXPECT_CALL(observer, OnFilesSelected).Times(0);
EXPECT_FALSE(UploadFileSync());
EXPECT_EQ(dialog_params_.type, ui::SelectFileDialog::SELECT_OPEN_MULTI_FILE);
// No extension filtering — all files are accepted, matching how Chromium's
// Browser::OpenFile works. The renderer handles MIME sniffing.
ASSERT_TRUE(dialog_params_.file_types);
ASSERT_EQ(1u, dialog_params_.file_types->extensions.size());
EXPECT_TRUE(std::ranges::contains(dialog_params_.file_types->extensions[0],
FILE_PATH_LITERAL("png")));
EXPECT_TRUE(std::ranges::contains(dialog_params_.file_types->extensions[0],
FILE_PATH_LITERAL("jpeg")));
EXPECT_TRUE(std::ranges::contains(dialog_params_.file_types->extensions[0],
FILE_PATH_LITERAL("jpg")));
EXPECT_TRUE(std::ranges::contains(dialog_params_.file_types->extensions[0],
FILE_PATH_LITERAL("webp")));
EXPECT_TRUE(std::ranges::contains(dialog_params_.file_types->extensions[0],
FILE_PATH_LITERAL("pdf")));
EXPECT_TRUE(dialog_params_.file_types->extensions.empty());
#if BUILDFLAG(IS_ANDROID)
// Android doesn't support view-source for text file extraction, so only
// image and PDF MIME types are accepted.
@@ -302,8 +294,13 @@ TEST_F(UploadFileHelperTest, PdfFileWithInvalidHeader) {
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
// Should fail since it has .pdf extension but doesn't look like a PDF
EXPECT_FALSE(result);
// Should return a stub entry (empty data, kText) so the frontend can
// detect the rejection and show an error.
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_TRUE((*result)[0]->data.empty());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
EXPECT_FALSE((*result)[0]->extracted_text.has_value());
}
TEST_F(UploadFileHelperTest, PdfFileTooSmall) {
@@ -324,8 +321,117 @@ TEST_F(UploadFileHelperTest, PdfFileTooSmall) {
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
// Should fail since it has .pdf extension but is too small to be a valid PDF
EXPECT_FALSE(result);
// Should return a stub entry for the rejected file.
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_TRUE((*result)[0]->data.empty());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
EXPECT_FALSE((*result)[0]->extracted_text.has_value());
}
TEST_F(UploadFileHelperTest, BinaryFileRejected) {
// Test that files with known binary MIME types are rejected
base::FilePath zip_path = temp_dir_.GetPath().AppendASCII("archive.zip");
ASSERT_TRUE(base::WriteFile(zip_path, "PK\x03\x04 fake zip content"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{zip_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
// Should return a stub entry for the rejected zip file.
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_TRUE((*result)[0]->data.empty());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
EXPECT_FALSE((*result)[0]->extracted_text.has_value());
}
TEST_F(UploadFileHelperTest, TextFileHandling) {
// Test text file with known extension
base::FilePath txt_path = temp_dir_.GetPath().AppendASCII("readme.txt");
ASSERT_TRUE(base::WriteFile(txt_path, "Hello, world!"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{txt_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_EQ((*result)[0]->filename, ExpectedFilename(txt_path));
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
EXPECT_EQ((*result)[0]->data.size(), 13u);
}
TEST_F(UploadFileHelperTest, TextFileWithoutExtension) {
// Test file without extension — should be treated as text
base::FilePath no_ext_path = temp_dir_.GetPath().AppendASCII("textfile");
ASSERT_TRUE(base::WriteFile(no_ext_path, "some text content"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{no_ext_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
}
TEST_F(UploadFileHelperTest, TextFileWithTrailingDot) {
// Test file ending with a dot — should be treated as text
base::FilePath dot_path = temp_dir_.GetPath().AppendASCII("file.");
ASSERT_TRUE(base::WriteFile(dot_path, "trailing dot content"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{dot_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
}
TEST_F(UploadFileHelperTest, TextFileWithUnknownExtension) {
// Test file with extension not in Chromium's MIME registry
base::FilePath diff_path = temp_dir_.GetPath().AppendASCII("changes.diff");
ASSERT_TRUE(base::WriteFile(diff_path, "--- a/file\n+++ b/file\n"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{diff_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
ASSERT_TRUE(result);
ASSERT_EQ(1u, result->size());
EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kText);
}
TEST_F(UploadFileHelperTest, MixedFileTypes) {
@@ -360,4 +466,42 @@ TEST_F(UploadFileHelperTest, MixedFileTypes) {
EXPECT_GT((*result)[1]->data.size(), 0u);
}
TEST_F(UploadFileHelperTest, MixedWithUnsupportedFileDropsInvalid) {
data_decoder::test::InProcessDataDecoder data_decoder;
// Upload a valid image and text file alongside an unsupported zip file.
auto png_bytes = gfx::test::CreatePNGBytes(100);
base::FilePath png_path = temp_dir_.GetPath().AppendASCII("photo.png");
base::FilePath txt_path = temp_dir_.GetPath().AppendASCII("readme.txt");
base::FilePath zip_path = temp_dir_.GetPath().AppendASCII("archive.zip");
ASSERT_TRUE(base::WriteFile(png_path, base::span(*png_bytes)));
ASSERT_TRUE(base::WriteFile(txt_path, "hello world"));
ASSERT_TRUE(base::WriteFile(zip_path, "PK\x03\x04 fake zip"));
ui::SelectFileDialog::SetFactory(
std::make_unique<content::FakeSelectFileDialogFactory>(
std::vector<base::FilePath>{zip_path, png_path, txt_path}));
testing::NiceMock<MockObserver> observer(file_helper_.get());
EXPECT_CALL(observer, OnFilesSelected).Times(1);
auto result = UploadFileSync();
testing::Mock::VerifyAndClearExpectations(&observer);
// All 3 files returned: image and text are valid, zip is a stub.
ASSERT_TRUE(result);
ASSERT_EQ(3u, result->size());
// Barrier callback order is nondeterministic, so check by type.
EXPECT_TRUE(std::ranges::any_of(*result, [](const auto& f) {
return f->type == mojom::UploadedFileType::kImage && !f->data.empty();
}));
// The text file and zip stub both have kText type. Distinguish by data.
EXPECT_TRUE(std::ranges::any_of(*result, [](const auto& f) {
return f->type == mojom::UploadedFileType::kText && !f->data.empty();
}));
// Zip stub: kText with empty data
EXPECT_TRUE(std::ranges::any_of(*result, [](const auto& f) {
return f->type == mojom::UploadedFileType::kText && f->data.empty();
}));
}
} // namespace ai_chat
@@ -239,46 +239,70 @@ void AIChatUIPageHandler::UploadFile(bool use_media_capture,
void AIChatUIPageHandler::OnFilesUploaded(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files) {
#if BUILDFLAG(ENABLE_PDF)
if (uploaded_files) {
// Collect PDF file paths and their indices before moving uploaded_files.
std::vector<std::pair<size_t, base::FilePath>> pdf_extractions;
// Collect all files needing text extraction in one pass.
// Store type alongside index/path so we can create the right extractor
// after uploaded_files is moved into the barrier callback.
struct ExtractionInfo {
size_t index;
base::FilePath path;
mojom::UploadedFileType type;
};
std::vector<ExtractionInfo> extractions;
for (size_t i = 0; i < uploaded_files->size(); ++i) {
auto& file = (*uploaded_files)[i];
if (file->type == mojom::UploadedFileType::kPdf &&
!file->extracted_text.has_value()) {
pdf_extractions.emplace_back(
i, base::FilePath::FromUTF8Unsafe(file->filename));
if (file->extracted_text.has_value() || file->data.empty()) {
continue;
}
bool needs_extraction = false;
#if BUILDFLAG(ENABLE_PDF)
needs_extraction |= file->type == mojom::UploadedFileType::kPdf;
#endif
#if !BUILDFLAG(IS_ANDROID)
needs_extraction |= file->type == mojom::UploadedFileType::kText;
#endif
if (needs_extraction) {
extractions.push_back(
{i, base::FilePath::FromUTF8Unsafe(file->filename), file->type});
}
}
if (!pdf_extractions.empty()) {
// Extract all PDFs in parallel via BarrierCallback.
if (!extractions.empty()) {
auto barrier =
base::BarrierCallback<std::pair<size_t, std::optional<std::string>>>(
pdf_extractions.size(),
base::BindOnce(&AIChatUIPageHandler::OnAllPdfTextsExtracted,
extractions.size(),
base::BindOnce(&AIChatUIPageHandler::OnAllFilesExtracted,
weak_ptr_factory_.GetWeakPtr(),
std::move(callback), std::move(uploaded_files)));
for (const auto& [idx, pdf_path] : pdf_extractions) {
auto extractor = std::make_unique<PdfTextExtractor>();
for (const auto& info : extractions) {
std::unique_ptr<FileTextExtractorBase> extractor;
#if BUILDFLAG(ENABLE_PDF)
if (info.type == mojom::UploadedFileType::kPdf) {
extractor = std::make_unique<PdfTextExtractor>();
}
#endif
#if !BUILDFLAG(IS_ANDROID)
if (info.type == mojom::UploadedFileType::kText) {
extractor = std::make_unique<TextFileExtractor>();
}
#endif
CHECK(extractor);
auto* extractor_ptr = extractor.get();
pdf_extractors_.push_back(std::move(extractor));
extractors_.push_back(std::move(extractor));
extractor_ptr->ExtractText(
profile_, pdf_path,
base::BindOnce(&AIChatUIPageHandler::OnSinglePdfTextExtracted,
weak_ptr_factory_.GetWeakPtr(), extractor_ptr, idx,
barrier));
profile_, info.path,
base::BindOnce(&AIChatUIPageHandler::OnSingleFileExtracted,
weak_ptr_factory_.GetWeakPtr(), extractor_ptr,
info.index, barrier));
}
return;
}
}
#endif // BUILDFLAG(ENABLE_PDF)
FinishUpload(std::move(callback), std::move(uploaded_files));
}
void AIChatUIPageHandler::OnAllPdfTextsExtracted(
void AIChatUIPageHandler::OnAllFilesExtracted(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files,
std::vector<std::pair<size_t, std::optional<std::string>>> results) {
@@ -292,6 +316,18 @@ void AIChatUIPageHandler::OnAllPdfTextsExtracted(
FinishUpload(std::move(callback), std::move(uploaded_files));
}
void AIChatUIPageHandler::OnSingleFileExtracted(
FileTextExtractorBase* extractor_ptr,
size_t file_index,
base::OnceCallback<void(std::pair<size_t, std::optional<std::string>>)>
barrier_cb,
std::optional<std::string> extracted_text) {
std::erase_if(extractors_, [extractor_ptr](const auto& e) {
return e.get() == extractor_ptr;
});
std::move(barrier_cb).Run({file_index, std::move(extracted_text)});
}
void AIChatUIPageHandler::FinishUpload(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files) {
@@ -327,20 +363,28 @@ void AIChatUIPageHandler::ProcessImageFile(
filename, std::move(callback)));
}
void AIChatUIPageHandler::ProcessTextFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessTextFileCallback callback) {
#if !BUILDFLAG(IS_ANDROID)
const auto extension = base::FilePath::FromUTF8Unsafe(filename).Extension();
ExtractAndProcessFile(
std::make_unique<TextFileExtractor>(), file_data,
extension.size() > 1 ? extension.substr(1) : FILE_PATH_LITERAL("txt"),
filename, mojom::UploadedFileType::kText, std::move(callback));
#else
// Android does not support background text extraction via view-source:.
std::move(callback).Run(nullptr);
#endif
}
void AIChatUIPageHandler::ProcessPdfFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessPdfFileCallback callback) {
#if BUILDFLAG(ENABLE_PDF)
auto extractor = std::make_unique<PdfTextExtractor>();
auto* extractor_ptr = extractor.get();
pdf_extractors_.push_back(std::move(extractor));
auto response_data = std::vector<uint8_t>(file_data);
extractor_ptr->ExtractText(
profile_, std::vector<uint8_t>(file_data),
base::BindOnce(&AIChatUIPageHandler::OnPdfTextExtracted,
weak_ptr_factory_.GetWeakPtr(), extractor_ptr, filename,
std::move(response_data), std::move(callback)));
ExtractAndProcessFile(std::make_unique<PdfTextExtractor>(), file_data,
FILE_PATH_LITERAL("pdf"), filename,
mojom::UploadedFileType::kPdf, std::move(callback));
#else
auto uploaded_file = ai_chat::mojom::UploadedFile::New(
filename, file_data.size(), file_data,
@@ -349,35 +393,46 @@ void AIChatUIPageHandler::ProcessPdfFile(const std::vector<uint8_t>& file_data,
#endif // BUILDFLAG(ENABLE_PDF)
}
#if BUILDFLAG(ENABLE_PDF)
void AIChatUIPageHandler::OnSinglePdfTextExtracted(
PdfTextExtractor* extractor_ptr,
size_t file_index,
base::OnceCallback<void(std::pair<size_t, std::optional<std::string>>)>
barrier_cb,
std::optional<std::string> extracted_text) {
std::erase_if(pdf_extractors_, [extractor_ptr](const auto& e) {
return e.get() == extractor_ptr;
});
std::move(barrier_cb).Run({file_index, std::move(extracted_text)});
void AIChatUIPageHandler::ExtractAndProcessFile(
std::unique_ptr<FileTextExtractorBase> extractor,
const std::vector<uint8_t>& file_data,
const base::FilePath::StringType& extension,
const std::string& filename,
mojom::UploadedFileType file_type,
base::OnceCallback<void(mojom::UploadedFilePtr)> callback) {
auto* extractor_ptr = extractor.get();
extractors_.push_back(std::move(extractor));
auto response_data = std::vector<uint8_t>(file_data);
extractor_ptr->ExtractText(
profile_, std::vector<uint8_t>(file_data), extension,
base::BindOnce(&AIChatUIPageHandler::OnFileExtracted,
weak_ptr_factory_.GetWeakPtr(), extractor_ptr, filename,
std::move(response_data), file_type, std::move(callback)));
}
void AIChatUIPageHandler::OnPdfTextExtracted(
PdfTextExtractor* extractor_ptr,
void AIChatUIPageHandler::OnFileExtracted(
FileTextExtractorBase* extractor_ptr,
std::string filename,
std::vector<uint8_t> file_data,
ProcessPdfFileCallback callback,
mojom::UploadedFileType file_type,
base::OnceCallback<void(mojom::UploadedFilePtr)> callback,
std::optional<std::string> extracted_text) {
auto file_size = file_data.size();
auto uploaded_file = ai_chat::mojom::UploadedFile::New(
std::move(filename), file_size, std::move(file_data),
ai_chat::mojom::UploadedFileType::kPdf, std::move(extracted_text));
std::move(callback).Run(std::move(uploaded_file));
std::erase_if(pdf_extractors_, [extractor_ptr](const auto& e) {
std::erase_if(extractors_, [extractor_ptr](const auto& e) {
return e.get() == extractor_ptr;
});
// Text extraction failure means the file is not text-renderable (e.g.
// binary). Return null so the frontend does not attach it.
if (!extracted_text.has_value() &&
file_type == mojom::UploadedFileType::kText) {
std::move(callback).Run(nullptr);
return;
}
auto file_size = file_data.size();
auto uploaded_file = ai_chat::mojom::UploadedFile::New(
std::move(filename), file_size, std::move(file_data), file_type,
std::move(extracted_text));
std::move(callback).Run(std::move(uploaded_file));
}
#endif // BUILDFLAG(ENABLE_PDF)
void AIChatUIPageHandler::GetPluralString(const std::string& key,
int32_t count,
@@ -14,9 +14,15 @@
#include "base/memory/weak_ptr.h"
#include "base/scoped_observation.h"
#include "base/task/cancelable_task_tracker.h"
#include "brave/browser/ai_chat/file_text_extractor_base.h"
#include "brave/browser/ai_chat/upload_file_helper.h"
#include "build/build_config.h"
#include "pdf/buildflags.h"
#if !BUILDFLAG(IS_ANDROID)
#include "brave/browser/ai_chat/text_file_extractor.h"
#endif
#if BUILDFLAG(ENABLE_PDF)
#include "brave/browser/ai_chat/pdf_text_extractor.h"
#endif
@@ -96,6 +102,9 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
void ProcessPdfFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessPdfFileCallback callback) override;
void ProcessTextFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessTextFileCallback callback) override;
void GetPluralString(const std::string& key,
int32_t count,
GetPluralStringCallback callback) override;
@@ -124,12 +133,16 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
private:
FRIEND_TEST_ALL_PREFIXES(AIChatUIPageHandlerBrowserTest,
OnFilesUploaded_WithPdf);
#if !BUILDFLAG(IS_ANDROID)
FRIEND_TEST_ALL_PREFIXES(AIChatUIPageHandlerBrowserTest,
OnFilesUploaded_WithText);
#endif
FRIEND_TEST_ALL_PREFIXES(AIChatUIPageHandlerTest,
FinishUpload_StripsPathToBasename);
FRIEND_TEST_ALL_PREFIXES(AIChatUIPageHandlerTest,
OnFilesUploaded_NonPdfGoesToFinish);
FRIEND_TEST_ALL_PREFIXES(AIChatUIPageHandlerTest,
OnAllPdfTextsExtracted_AppliesResults);
OnAllFilesExtracted_AppliesResults);
class ChatContextObserver : public content::WebContentsObserver {
public:
@@ -165,26 +178,37 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
void OnFilesUploaded(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files);
void OnAllPdfTextsExtracted(
// Collects barrier results and applies extracted text to uploaded files.
void OnAllFilesExtracted(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files,
std::vector<std::pair<size_t, std::optional<std::string>>> results);
void FinishUpload(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files);
#if BUILDFLAG(ENABLE_PDF)
void OnSinglePdfTextExtracted(
PdfTextExtractor* extractor_ptr,
// Per-file barrier callback — removes the finished extractor and forwards.
void OnSingleFileExtracted(
FileTextExtractorBase* extractor_ptr,
size_t file_index,
base::OnceCallback<void(std::pair<size_t, std::optional<std::string>>)>
barrier_cb,
std::optional<std::string> extracted_text);
void OnPdfTextExtracted(PdfTextExtractor* extractor_ptr,
std::string filename,
std::vector<uint8_t> file_data,
ProcessPdfFileCallback callback,
std::optional<std::string> extracted_text);
#endif // BUILDFLAG(ENABLE_PDF)
void FinishUpload(
UploadFileCallback callback,
std::optional<std::vector<mojom::UploadedFilePtr>> uploaded_files);
// Shared helper for ProcessTextFile / ProcessPdfFile.
void ExtractAndProcessFile(
std::unique_ptr<FileTextExtractorBase> extractor,
const std::vector<uint8_t>& file_data,
const base::FilePath::StringType& extension,
const std::string& filename,
mojom::UploadedFileType file_type,
base::OnceCallback<void(mojom::UploadedFilePtr)> callback);
void OnFileExtracted(
FileTextExtractorBase* extractor_ptr,
std::string filename,
std::vector<uint8_t> file_data,
mojom::UploadedFileType file_type,
base::OnceCallback<void(mojom::UploadedFilePtr)> callback,
std::optional<std::string> extracted_text);
raw_ptr<AIChatTabHelper> active_chat_tab_helper_ = nullptr;
raw_ptr<content::WebContents> owner_web_contents_ = nullptr;
@@ -203,10 +227,8 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
// DataDecoder instance for processing image data
data_decoder::DataDecoder data_decoder_;
#if BUILDFLAG(ENABLE_PDF)
// Active PDF text extractors (owned until extraction completes)
std::vector<std::unique_ptr<PdfTextExtractor>> pdf_extractors_;
#endif
// Active file extractors (owned until extraction completes)
std::vector<std::unique_ptr<FileTextExtractorBase>> extractors_;
mojo::Receiver<ai_chat::mojom::AIChatUIHandler> receiver_;
mojo::Remote<ai_chat::mojom::ChatUI> chat_ui_;
@@ -24,6 +24,7 @@
#include "brave/components/ai_chat/core/common/features.h"
#include "brave/components/ai_chat/core/common/mojom/tab_tracker.mojom.h"
#include "brave/components/constants/brave_paths.h"
#include "build/build_config.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/test/base/chrome_test_utils.h"
#include "chrome/test/base/in_process_browser_test.h"
@@ -47,6 +48,11 @@
namespace ai_chat {
#if !BUILDFLAG(IS_ANDROID)
constexpr char kExpectedTextContent[] =
"Hello from a text file.\nThis is line two.";
#endif
#if BUILDFLAG(ENABLE_PDF)
constexpr char kExpectedPdfText[] = "This is the way\nI have spoken";
#endif // BUILDFLAG(ENABLE_PDF)
@@ -339,6 +345,103 @@ IN_PROC_BROWSER_TEST_F(AIChatUIPageHandlerBrowserTest,
}
#endif // BUILDFLAG(ENABLE_PDF)
#if !BUILDFLAG(IS_ANDROID)
IN_PROC_BROWSER_TEST_F(AIChatUIPageHandlerBrowserTest, ProcessTextFile) {
auto* page_handler = GetPageHandler(web_contents());
ASSERT_TRUE(page_handler);
base::FilePath txt_path;
std::vector<uint8_t> txt_bytes;
{
base::ScopedAllowBlockingForTesting allow_blocking;
txt_path = base::PathService::CheckedGet(brave::DIR_TEST_DATA)
.AppendASCII("leo")
.AppendASCII("dummy.txt");
auto bytes = base::ReadFileToBytes(txt_path);
ASSERT_TRUE(bytes.has_value());
txt_bytes = std::move(*bytes);
}
base::test::TestFuture<ai_chat::mojom::UploadedFilePtr> future;
page_handler->ProcessTextFile(txt_bytes, "dummy.txt", future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result);
EXPECT_EQ(result->filename, "dummy.txt");
EXPECT_EQ(result->type, ai_chat::mojom::UploadedFileType::kText);
EXPECT_GT(result->data.size(), 0u);
ASSERT_TRUE(result->extracted_text.has_value());
EXPECT_EQ(*result->extracted_text, kExpectedTextContent);
}
// Verify that HTML files are not rendered (no script execution, no external
// resource loading). The extracted text must be the raw HTML source.
IN_PROC_BROWSER_TEST_F(AIChatUIPageHandlerBrowserTest, ProcessHtmlFile) {
auto* page_handler = GetPageHandler(web_contents());
ASSERT_TRUE(page_handler);
base::FilePath html_path;
std::vector<uint8_t> html_bytes;
{
base::ScopedAllowBlockingForTesting allow_blocking;
html_path = base::PathService::CheckedGet(brave::DIR_TEST_DATA)
.AppendASCII("leo")
.AppendASCII("dummy.html");
auto bytes = base::ReadFileToBytes(html_path);
ASSERT_TRUE(bytes.has_value());
html_bytes = std::move(*bytes);
}
base::test::TestFuture<ai_chat::mojom::UploadedFilePtr> future;
page_handler->ProcessTextFile(html_bytes, "dummy.html", future.GetCallback());
auto result = future.Take();
ASSERT_TRUE(result);
EXPECT_EQ(result->filename, "dummy.html");
ASSERT_TRUE(result->extracted_text.has_value());
// Raw source must contain HTML tags — proves it was NOT rendered.
EXPECT_TRUE(result->extracted_text->find("<p>Hello from an HTML file.</p>") !=
std::string::npos);
EXPECT_TRUE(result->extracted_text->find("<script>") != std::string::npos);
}
IN_PROC_BROWSER_TEST_F(AIChatUIPageHandlerBrowserTest,
OnFilesUploaded_WithText) {
auto* page_handler = GetPageHandler(web_contents());
ASSERT_TRUE(page_handler);
base::FilePath txt_path;
std::vector<uint8_t> txt_bytes;
{
base::ScopedAllowBlockingForTesting allow_blocking;
txt_path = base::PathService::CheckedGet(brave::DIR_TEST_DATA)
.AppendASCII("leo")
.AppendASCII("dummy.txt");
auto bytes = base::ReadFileToBytes(txt_path);
ASSERT_TRUE(bytes.has_value());
txt_bytes = std::move(*bytes);
}
std::vector<ai_chat::mojom::UploadedFilePtr> files;
files.push_back(ai_chat::mojom::UploadedFile::New(
txt_path.AsUTF8Unsafe(), txt_bytes.size(), std::move(txt_bytes),
ai_chat::mojom::UploadedFileType::kText, std::nullopt));
base::test::TestFuture<
std::optional<std::vector<ai_chat::mojom::UploadedFilePtr>>>
future;
page_handler->OnFilesUploaded(future.GetCallback(),
std::make_optional(std::move(files)));
auto result = future.Take();
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->filename, "dummy.txt");
ASSERT_TRUE((*result)[0]->extracted_text.has_value());
EXPECT_EQ(*(*result)[0]->extracted_text, kExpectedTextContent);
}
#endif // !BUILDFLAG(IS_ANDROID)
IN_PROC_BROWSER_TEST_F(AIChatUIPageHandlerBrowserTest,
AssociateURLDoesNotCrashShutdown) {
auto* ai_chat_contents = web_contents();
@@ -166,7 +166,7 @@ TEST_F(AIChatUIPageHandlerTest, OnFilesUploaded_NonPdfGoesToFinish) {
EXPECT_EQ((*result)[0]->filename, "photo.png");
}
TEST_F(AIChatUIPageHandlerTest, OnAllPdfTextsExtracted_AppliesResults) {
TEST_F(AIChatUIPageHandlerTest, OnAllFilesExtracted_AppliesResults) {
std::vector<mojom::UploadedFilePtr> files;
files.push_back(
mojom::UploadedFile::New("/path/doc1.pdf", 100, std::vector<uint8_t>(100),
@@ -177,30 +177,45 @@ TEST_F(AIChatUIPageHandlerTest, OnAllPdfTextsExtracted_AppliesResults) {
files.push_back(
mojom::UploadedFile::New("/path/doc2.pdf", 200, std::vector<uint8_t>(200),
mojom::UploadedFileType::kPdf, std::nullopt));
files.push_back(mojom::UploadedFile::New(
"/path/config.conf", 80, std::vector<uint8_t>(80),
mojom::UploadedFileType::kText, std::nullopt));
files.push_back(mojom::UploadedFile::New(
"/path/script.py", 120, std::vector<uint8_t>(120),
mojom::UploadedFileType::kText, std::nullopt));
std::vector<std::pair<size_t, std::optional<std::string>>> results;
results.emplace_back(0, "Text from doc1");
results.emplace_back(2, std::nullopt); // extraction failed for doc2
results.emplace_back(3, "key=value");
results.emplace_back(4, std::nullopt); // extraction failed for script.py
base::test::TestFuture<std::optional<std::vector<mojom::UploadedFilePtr>>>
future;
page_handler()->OnAllPdfTextsExtracted(future.GetCallback(),
std::make_optional(std::move(files)),
std::move(results));
page_handler()->OnAllFilesExtracted(future.GetCallback(),
std::make_optional(std::move(files)),
std::move(results));
auto result = future.Take();
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result->size(), 3u);
ASSERT_EQ(result->size(), 5u);
// PDF extracted text applied and path stripped
EXPECT_EQ((*result)[0]->filename, "doc1.pdf");
ASSERT_TRUE((*result)[0]->extracted_text.has_value());
EXPECT_EQ(*(*result)[0]->extracted_text, "Text from doc1");
// Non-PDF unaffected, path stripped
// Image unaffected, path stripped
EXPECT_EQ((*result)[1]->filename, "photo.png");
EXPECT_FALSE((*result)[1]->extracted_text.has_value());
// Failed extraction, path stripped
// Failed PDF extraction, path stripped
EXPECT_EQ((*result)[2]->filename, "doc2.pdf");
EXPECT_FALSE((*result)[2]->extracted_text.has_value());
// Text extracted and path stripped
EXPECT_EQ((*result)[3]->filename, "config.conf");
ASSERT_TRUE((*result)[3]->extracted_text.has_value());
EXPECT_EQ(*(*result)[3]->extracted_text, "key=value");
// Failed text extraction, path stripped
EXPECT_EQ((*result)[4]->filename, "script.py");
EXPECT_FALSE((*result)[4]->extracted_text.has_value());
}
} // namespace ai_chat
@@ -310,6 +310,7 @@ std::vector<OAIMessage> BuildOAIMessages(
std::vector<mojom::ContentBlockPtr> uploaded_images_content_blocks;
std::vector<mojom::ContentBlockPtr> screenshots_content_blocks;
std::vector<mojom::ContentBlockPtr> uploaded_pdfs_content_blocks;
std::vector<mojom::ContentBlockPtr> uploaded_text_files_content_blocks;
uploaded_images_content_blocks.push_back(
mojom::ContentBlock::NewTextContentBlock(mojom::TextContentBlock::New(
@@ -320,6 +321,9 @@ std::vector<OAIMessage> BuildOAIMessages(
uploaded_pdfs_content_blocks.push_back(
mojom::ContentBlock::NewTextContentBlock(mojom::TextContentBlock::New(
"These PDFs are uploaded by the user")));
uploaded_text_files_content_blocks.push_back(
mojom::ContentBlock::NewTextContentBlock(mojom::TextContentBlock::New(
"These text files are uploaded by the user")));
for (const auto& uploaded_file : *message->uploaded_files) {
if (uploaded_file->type == mojom::UploadedFileType::kImage ||
@@ -354,6 +358,18 @@ std::vector<OAIMessage> BuildOAIMessages(
EngineConsumer::GetPdfDataURL(uploaded_file->data)),
pdf_filename)));
}
} else if (uploaded_file->type == mojom::UploadedFileType::kText) {
std::string text_filename = uploaded_file->filename.empty()
? "uploaded.txt"
: uploaded_file->filename;
if (uploaded_file->extracted_text.has_value() &&
!uploaded_file->extracted_text->empty()) {
uploaded_text_files_content_blocks.push_back(
mojom::ContentBlock::NewTextContentBlock(
mojom::TextContentBlock::New(
"[File: " + text_filename + "]\n" +
*uploaded_file->extracted_text)));
}
}
}
@@ -377,6 +393,13 @@ std::vector<OAIMessage> BuildOAIMessages(
std::make_move_iterator(uploaded_pdfs_content_blocks.begin()),
std::make_move_iterator(uploaded_pdfs_content_blocks.end()));
}
if (uploaded_text_files_content_blocks.size() > 1) {
oai_message.content.insert(
oai_message.content.end(),
std::make_move_iterator(uploaded_text_files_content_blocks.begin()),
std::make_move_iterator(uploaded_text_files_content_blocks.end()));
}
}
if (message->selected_text.has_value() &&
@@ -585,6 +585,60 @@ TEST_F(OAIMessageUtilsTest, BuildOAIMessages_PdfExtractedTextPreferred) {
VerifyTextBlock(FROM_HERE, messages[0].content[3], "query0");
}
TEST_F(OAIMessageUtilsTest, BuildOAIMessages_TextFileExtractedText) {
// Test that text files with extracted_text produce TextContentBlock with
// [File: filename] prefix, and text files without extracted_text are skipped.
auto history = CreateSampleChatHistory(1);
auto text_files =
CreateSampleUploadedFiles(2, mojom::UploadedFileType::kText);
// First text file has extracted text — should produce TextContentBlock
text_files[0]->extracted_text = "config_key=config_value";
text_files[0]->filename = "app.conf";
// Second text file has no extracted text — should be skipped
text_files[1]->filename = "failed.txt";
history[0]->uploaded_files = std::move(text_files);
PageContentsMap page_contents_map;
std::vector<OAIMessage> messages =
BuildOAIMessages(std::move(page_contents_map), history, nullptr, true,
10000, [](std::string&) {});
ASSERT_EQ(messages.size(), 2u);
EXPECT_EQ(messages[0].role, "user");
// Content: text files intro + extracted text + prompt = 3 blocks
// (second file with no extracted_text is skipped)
ASSERT_EQ(messages[0].content.size(), 3u);
VerifyTextBlock(FROM_HERE, messages[0].content[0],
"These text files are uploaded by the user");
VerifyTextBlock(FROM_HERE, messages[0].content[1],
"[File: app.conf]\nconfig_key=config_value");
VerifyTextBlock(FROM_HERE, messages[0].content[2], "query0");
}
TEST_F(OAIMessageUtilsTest, BuildOAIMessages_TextFileDefaultFilename) {
// Test that text files with empty filename get default "uploaded.txt"
auto history = CreateSampleChatHistory(1);
auto text_files =
CreateSampleUploadedFiles(1, mojom::UploadedFileType::kText);
text_files[0]->extracted_text = "some content";
text_files[0]->filename.clear();
history[0]->uploaded_files = std::move(text_files);
PageContentsMap page_contents_map;
std::vector<OAIMessage> messages =
BuildOAIMessages(std::move(page_contents_map), history, nullptr, true,
10000, [](std::string&) {});
ASSERT_EQ(messages.size(), 2u);
ASSERT_EQ(messages[0].content.size(), 3u);
VerifyTextBlock(FROM_HERE, messages[0].content[1],
"[File: uploaded.txt]\nsome content");
}
TEST_F(OAIMessageUtilsTest, BuildOAIMessages_Memory_Excluded) {
// Enable customization and set data
prefs_.SetBoolean(prefs::kBraveAIChatUserCustomizationEnabled, true);
@@ -218,6 +218,11 @@ interface AIChatUIHandler {
ProcessPdfFile(array<uint8> file_data, string filename)
=> (UploadedFile? processed_file);
// Process text files: extract text via background renderer for encoding
// detection
ProcessTextFile(array<uint8> file_data, string filename)
=> (UploadedFile? processed_file);
// Get a plural string for the given key and count.
GetPluralString(string key, int32 count) => (string plural_string);
@@ -113,6 +113,7 @@ enum UploadedFileType {
kImage = 0,
kScreenshot,
kPdf,
kText,
};
// This does not cover more specific data that the Service owns, such as the
@@ -151,9 +152,9 @@ struct UploadedFile {
uint32 filesize;
array<uint8> data;
UploadedFileType type;
// Text extracted from PDF files via the background PDF viewer and ScreenAI
// OCR pipeline. Null for non-PDF files or when extraction failed/timed out.
// When present, the engine layer sends this text instead of raw PDF bytes.
// Text extracted from PDF or text files via a background renderer. Null when
// extraction failed/timed out or not applicable. When present, the engine
// layer sends this text instead of raw file bytes.
string? extracted_text;
};
@@ -237,7 +237,8 @@ export const processUploadedFilesWithLimits = (
const isImage =
file.type === Mojom.UploadedFileType.kImage
|| file.type === Mojom.UploadedFileType.kScreenshot
const isDocument = file.type === Mojom.UploadedFileType.kPdf
const isPdf = file.type === Mojom.UploadedFileType.kPdf
const isText = file.type === Mojom.UploadedFileType.kText
if (isImage) {
const maxNewImages =
Mojom.MAX_IMAGES - totalUploadedImages - currentPendingImages
@@ -245,7 +246,7 @@ export const processUploadedFilesWithLimits = (
newFiles.push(file)
currentImages++
}
} else if (isDocument) {
} else if (isPdf) {
const hasExtractedText = !!file.extractedText
if (hasExtractedText) {
// PDFs with extracted text bypass raw file limits
@@ -265,6 +266,9 @@ export const processUploadedFilesWithLimits = (
currentRawDocuments++
}
}
} else if (isText) {
// Text files are not subject to document count/size limits
newFiles.push(file)
}
}
@@ -104,6 +104,9 @@ export default function createAIChatApi(
processPdfFile: {
mutationResponse: (result) => result.processedFile,
},
processTextFile: {
mutationResponse: (result) => result.processedFile,
},
getPluralString: {
response: (result) => result.pluralString,
},
@@ -213,6 +213,16 @@ export function createMockUIHandler(
extractedText: undefined,
},
}),
processTextFile: () =>
Promise.resolve({
processedFile: {
filename: '',
filesize: 0,
data: [],
type: Mojom.UploadedFileType.kText,
extractedText: undefined,
},
}),
getPluralString: () => Promise.resolve({ pluralString: '' }),
setChatUI: () => Promise.resolve({ isStandalone: false }),
@@ -175,6 +175,7 @@ function AttachmentUploadItem({
file.type === Mojom.UploadedFileType.kImage
|| file.type === Mojom.UploadedFileType.kScreenshot
const isPdf = file.type === Mojom.UploadedFileType.kPdf
const isText = file.type === Mojom.UploadedFileType.kText
const isFileFullPageScreenshot = isFullPageScreenshot(file)
const dataUrl = React.useMemo(() => {
@@ -208,7 +209,7 @@ function AttachmentUploadItem({
className={className}
/>
)
} else if (isPdf) {
} else if (isPdf || isText) {
return (
<AttachmentItem
icon={<Icon name='file' />}
@@ -5,9 +5,10 @@
import * as React from 'react'
import { getLocale } from '$web-common/locale'
import { showAlert } from '@brave/leo/react/alertCenter'
import Icon from '@brave/leo/react/icon'
import * as Mojom from '../../../common/mojom'
import styles from './style.module.scss'
import { isImageFile, isPdfFile } from '../../constants/file_types'
import { useAIChat } from '../../state/ai_chat_context'
import { useConversation } from '../../state/conversation_context'
import { convertFileToUploadedFile } from '../../utils/file_utils'
@@ -21,25 +22,36 @@ export default function DragOverlay() {
e.stopPropagation()
clearDragState()
const files = Array.from(e.dataTransfer?.files || []).filter(
(file) => isImageFile(file) || isPdfFile(file),
)
const files = Array.from(e.dataTransfer?.files || [])
if (files.length === 0) {
return
}
try {
const uploadedFiles = await Promise.all(
const results = await Promise.all(
files.map((file) =>
convertFileToUploadedFile(
file,
aiChat.processImageFile,
aiChat.processPdfFile,
aiChat.processTextFile,
),
),
)
attachImages(uploadedFiles)
const uploadedFiles = results.filter(
(f): f is Mojom.UploadedFile => f !== null,
)
if (uploadedFiles.length > 0) {
attachImages(uploadedFiles)
}
if (uploadedFiles.length < files.length) {
showAlert({
type: 'error',
content: getLocale(S.CHAT_UI_FILE_UPLOAD_ERROR),
actions: [],
})
}
} catch (error) {
// Silently fail - error will be handled by the upload system
}
@@ -436,16 +436,23 @@ describe('input box', () => {
jest.clearAllMocks()
mockConvertFileToUploadedFile.mockImplementation((file: File) => {
const mimeType = file.type.toLowerCase()
let type = UploadedFileType.kText
if (mimeType.startsWith('image/')) {
type = UploadedFileType.kImage
} else if (mimeType === 'application/pdf') {
type = UploadedFileType.kPdf
}
return Promise.resolve({
filename: file.name,
filesize: file.size,
data: Array.from(new Uint8Array(8)), // Mock data array
type: UploadedFileType.kImage,
data: Array.from(new Uint8Array(8)),
type,
})
})
})
it('filters image files and calls attachImages on paste', async () => {
it('accepts all file types on paste', async () => {
const mockAttachImages = jest.fn()
const { container } = await renderInputBox(
<MockContext>
@@ -488,6 +495,12 @@ describe('input box', () => {
data: expect.any(Array),
type: UploadedFileType.kImage,
}),
expect.objectContaining({
filename: 'test.txt',
filesize: 1024,
data: expect.any(Array),
type: UploadedFileType.kText,
}),
])
})
})
@@ -3,6 +3,7 @@
* 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 { showAlert } from '@brave/leo/react/alertCenter'
import Icon from '@brave/leo/react/icon'
import Button from '@brave/leo/react/button'
import Tooltip from '@brave/leo/react/tooltip'
@@ -23,7 +24,6 @@ import {
} from '../attachment_item'
import { ModelSelector } from '../model_selector'
import usePromise from '$web-common/usePromise'
import { isImageFile } from '../../constants/file_types'
import { convertFileToUploadedFile } from '../../utils/file_utils'
import { isFullPageScreenshot } from '../../../common/conversation_history_utils'
import Editable from './editable'
@@ -76,6 +76,7 @@ type Props = Pick<
| 'getPluralString'
| 'processImageFile'
| 'processPdfFile'
| 'processTextFile'
| 'openAIChatAgentProfile'
| 'skills'
>
@@ -219,26 +220,39 @@ function InputBox(props: InputBoxProps) {
return
}
const files = Array.from(clipboardData.files).filter(isImageFile)
const files = Array.from(clipboardData.files)
if (files.length === 0) {
return
}
// Prevent the default paste behavior for images
// Prevent the default paste behavior for files
e.preventDefault()
try {
const uploadedFiles = await Promise.all(
const results = await Promise.all(
files.map((file) =>
convertFileToUploadedFile(
file,
props.context.processImageFile,
props.context.processPdfFile,
props.context.processTextFile,
),
),
)
props.context.attachImages(uploadedFiles)
const uploadedFiles = results.filter(
(f): f is Mojom.UploadedFile => f !== null,
)
if (uploadedFiles.length > 0) {
props.context.attachImages(uploadedFiles)
}
if (uploadedFiles.length < files.length) {
showAlert({
type: 'error',
content: getLocale(S.CHAT_UI_FILE_UPLOAD_ERROR),
actions: [],
})
}
} catch (error) {
// Silently fail - error will be handled by the upload system
}
@@ -152,6 +152,7 @@ export default function useProvideAIChatContext(props: AIChatContextProps) {
// and we can do that via monitoring the mutation in the provided hook.
processImageFile: api.processImageFile,
processPdfFile: api.processPdfFile,
processTextFile: api.processTextFile,
/**
* @deprecated use api.uiHandler.openAIChatAgentProfile directly instead
@@ -4,7 +4,9 @@
// You can obtain one at https://mozilla.org/MPL/2.0/.
import * as React from 'react'
import { showAlert } from '@brave/leo/react/alertCenter'
import generateReactContext from '$web-common/api/react_api'
import { getLocale } from '$web-common/locale'
import { Url } from 'gen/url/mojom/url.mojom.m.js'
import { IGNORE_EXTERNAL_LINK_WARNING_KEY } from '../../common/constants'
import {
@@ -391,6 +393,24 @@ export function useProvideConversationContext(props: ConversationContextProps) {
}
const processUploadedFiles = async (files: Mojom.UploadedFile[]) => {
// Filter out text files where extraction failed (no extracted text).
const validFiles = files.filter(
(f) =>
f.type !== Mojom.UploadedFileType.kText
|| (f.extractedText !== undefined && f.extractedText !== null),
)
// Show error when some files were dropped: either text extraction
// failed, or unsupported files were included (the backend returns
// empty stubs for unsupported types like zip so they are filtered
// out here, while cancellation returns null and skips this path).
if (validFiles.length < files.length) {
showAlert({
type: 'error',
content: getLocale(S.CHAT_UI_FILE_UPLOAD_ERROR),
actions: [],
})
}
// After mutation, any returned promise will be awaited before settling.
// This won't re-fetch the conversation history, just get the latest
// version if it's not invalidated.
@@ -399,7 +419,7 @@ export function useProvideConversationContext(props: ConversationContextProps) {
// data.
setPendingMessageFiles((pendingMessageFiles) => {
const newFiles = processUploadedFilesWithLimits(
files,
validFiles,
conversationHistory,
pendingMessageFiles,
)
@@ -433,7 +453,8 @@ export function useProvideConversationContext(props: ConversationContextProps) {
) => {
uploadFileMutation.mutate(args, {
onSuccess: async (uploadedFiles, [useMediaCapture]) => {
// Reset event state, avoid us having to make a useState<bool> for this
// Reset event state, avoid us having to make a useState<bool>
// for this
aiChat.api.resetOnUploadFilesSelected()
if (uploadedFiles) {
return processUploadedFiles(uploadedFiles)
@@ -7,7 +7,6 @@ import {
convertFileToUploadedFile,
FileReadError,
ImageProcessingError,
UnsupportedFileTypeError,
} from './file_utils'
import * as Mojom from '../../common/mojom'
import type { AIChatContext } from '../state/ai_chat_context'
@@ -241,12 +240,18 @@ describe('convertFileToUploadedFile', () => {
)
})
it('throws UnsupportedFileTypeError for unknown file types', async () => {
it('treats text files as kText via processTextFile', async () => {
const file = createMockFile('test.txt', 'text/plain')
const mockArrayBuffer = new ArrayBuffer(8)
const expectedData = Array.from(new Uint8Array(mockArrayBuffer))
const mockProcessTextFile = jest.fn().mockResolvedValue({
filename: 'test.txt',
filesize: file.size,
data: expectedData,
type: Mojom.UploadedFileType.kText,
extractedText: 'extracted content',
})
// Override readAsArrayBuffer to trigger success
// (FileReader will work, but file type check will fail)
mockFileReader.readAsArrayBuffer.mockImplementation(() => {
process.nextTick(() => {
if (mockFileReader.onload) {
@@ -255,16 +260,45 @@ describe('convertFileToUploadedFile', () => {
})
})
await expect(
convertFileToUploadedFile(file, mockProcessImageFile),
).rejects.toThrow(UnsupportedFileTypeError)
await expect(
convertFileToUploadedFile(file, mockProcessImageFile),
).rejects.toThrow(
'Unsupported file type: text/plain. Only images and PDF files are '
+ 'supported.',
const result = await convertFileToUploadedFile(
file,
mockProcessImageFile,
undefined,
mockProcessTextFile,
)
expect(mockProcessImageFile).not.toHaveBeenCalled()
expect(mockProcessTextFile).toHaveBeenCalled()
expect(result).toEqual(
expect.objectContaining({
filename: 'test.txt',
type: Mojom.UploadedFileType.kText,
extractedText: 'extracted content',
}),
)
})
it('returns null for text files when extraction fails', async () => {
const file = createMockFile('binary.dat', 'application/octet-stream')
const mockArrayBuffer = new ArrayBuffer(8)
const mockProcessTextFile = jest.fn().mockResolvedValue(null)
mockFileReader.readAsArrayBuffer.mockImplementation(() => {
process.nextTick(() => {
if (mockFileReader.onload) {
mockFileReader.onload({ target: { result: mockArrayBuffer } })
}
})
})
const result = await convertFileToUploadedFile(
file,
mockProcessImageFile,
undefined,
mockProcessTextFile,
)
expect(result).toBeNull()
})
})
@@ -27,19 +27,13 @@ export class ImageProcessingError extends Error {
}
}
export class UnsupportedFileTypeError extends Error {
constructor(message: string) {
super(message)
this.name = 'UnsupportedFileTypeError'
}
}
// Utility function to convert File objects to UploadedFile format
export const convertFileToUploadedFile = async (
file: File,
processImageFile: AIChatContext['processImageFile'],
processPdfFile?: AIChatContext['processPdfFile'],
): Promise<Mojom.UploadedFile> => {
processTextFile?: AIChatContext['processTextFile'],
): Promise<Mojom.UploadedFile | null> => {
const reader = new FileReader()
const arrayBuffer = await new Promise<ArrayBuffer>((resolve, reject) => {
reader.onload = (e) => {
@@ -70,41 +64,41 @@ export const convertFileToUploadedFile = async (
}
}
// Fallback: return raw PDF data without extracted text
const uploadedFile: Mojom.UploadedFile = {
return {
filename: file.name,
filesize: file.size,
data: Array.from(uint8Array),
type: Mojom.UploadedFileType.kPdf,
extractedText: undefined,
}
return uploadedFile
} else if (mimeType.startsWith('image/')) {
// Use backend processing for images via mojo call
try {
const processedFile = await processImageFile([
const processedFile = await processImageFile([
Array.from(uint8Array),
file.name,
])
if (!processedFile) {
throw new ImageProcessingError(
'Failed to process image file: Backend returned no result',
)
}
return processedFile
} else {
// Everything else is treated as a text file. The renderer handles
// MIME sniffing and will render the file as text if possible.
// If extraction fails (e.g. binary file), null is returned and
// the file is not attached.
if (processTextFile) {
const processedFile = await processTextFile([
Array.from(uint8Array),
file.name,
])
if (!processedFile) {
throw new ImageProcessingError(
'Failed to process image file: Backend returned no result',
)
if (processedFile) {
return processedFile
}
return processedFile
} catch (error) {
if (error instanceof ImageProcessingError) {
throw error
}
// Re-throw any other errors as-is
throw error
}
} else {
throw new UnsupportedFileTypeError(
`Unsupported file type: ${file.type}. Only images and PDF files `
+ `are supported.`,
)
return null
}
}
@@ -45,6 +45,9 @@
<message name="IDS_CHAT_UI_DROP_FILES_DESCRIPTION" desc="Descriptive text shown when dragging files over the page" formatter_data="webui=AiChat">
Drop any files here to add them to the conversation
</message>
<message name="IDS_CHAT_UI_FILE_UPLOAD_ERROR" desc="Alert shown when one or more files could not be read as text and were not attached" formatter_data="webui=AiChat">
Some files could not be read and were not attached
</message>
<message name="IDS_CHAT_UI_SUGGEST_QUESTIONS_LABEL" desc="Button label to ask the AI Chat to suggest some related questions" formatter_data="webui=AiChat">
Suggest questions…
</message>
@@ -66,6 +66,9 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
void ProcessPdfFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessPdfFileCallback callback) override;
void ProcessTextFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessTextFileCallback callback) override;
void UploadFile(bool use_media_capture, UploadFileCallback callback) override;
void GetPluralString(const std::string& key,
int32_t count,
@@ -225,6 +225,17 @@ void AIChatUIPageHandler::ProcessPdfFile(const std::vector<uint8_t>& file_data,
std::move(callback).Run(std::move(uploaded_file));
}
void AIChatUIPageHandler::ProcessTextFile(const std::vector<uint8_t>& file_data,
const std::string& filename,
ProcessTextFileCallback callback) {
// iOS does not support background text extraction.
// Return the raw data without extracted text.
auto uploaded_file = ai_chat::mojom::UploadedFile::New(
filename, file_data.size(), file_data,
ai_chat::mojom::UploadedFileType::kText, std::nullopt);
std::move(callback).Run(std::move(uploaded_file));
}
void AIChatUIPageHandler::UploadFile(bool use_media_capture,
UploadFileCallback callback) {
id<AIChatUIHandlerBridge> bridge =
+9
View File
@@ -0,0 +1,9 @@
<!DOCTYPE html>
<html>
<head><title>Test</title></head>
<body>
<p>Hello from an HTML file.</p>
<img src="https://example.com/track.png">
<script>document.title = "executed";</script>
</body>
</html>
+2
View File
@@ -0,0 +1,2 @@
Hello from a text file.
This is line two.