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:
@@ -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) {
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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>
|
||||
@@ -0,0 +1,2 @@
|
||||
Hello from a text file.
|
||||
This is line two.
|
||||
Reference in New Issue
Block a user