From 16e68d6a9032ef60a569c3b508a801ce58a8eee9 Mon Sep 17 00:00:00 2001
From: Pete Miller <miller.pete@gmail.com>
Date: Wed, 11 Dec 2024 17:01:14 -0800
Subject: [PATCH] Don't crash when editing the title of a conversation not
 in-memory

---
 .../ai_chat/core/browser/ai_chat_service.cc   | 32 +++++++++----------
 .../ai_chat/core/browser/ai_chat_service.h    |  7 ++--
 .../core/browser/conversation_handler.cc      |  2 +-
 .../core/browser/conversation_handler.h       |  5 +--
 4 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/components/ai_chat/core/browser/ai_chat_service.cc b/components/ai_chat/core/browser/ai_chat_service.cc
index e45a02b043c8..4613b1b4d61a 100644
--- a/components/ai_chat/core/browser/ai_chat_service.cc
+++ b/components/ai_chat/core/browser/ai_chat_service.cc
@@ -61,7 +61,8 @@ constexpr auto kAllowedSchemes = base::MakeFixedFlatSet<std::string_view>(
     {url::kHttpsScheme, url::kHttpScheme, url::kFileScheme, url::kDataScheme});
 
 std::vector<mojom::Conversation*> FilterVisibleConversations(
-    std::map<std::string, mojom::ConversationPtr>& conversations_map) {
+    std::map<std::string, mojom::ConversationPtr, std::less<>>&
+        conversations_map) {
   std::vector<mojom::Conversation*> conversations;
   for (const auto& kv : conversations_map) {
     auto& conversation = kv.second;
@@ -598,14 +599,8 @@ void AIChatService::DeleteConversation(const std::string& id) {
 
 void AIChatService::RenameConversation(const std::string& id,
                                        const std::string& new_name) {
-  ConversationHandler* conversation_handler =
-      conversation_handlers_.at(id).get();
-  if (!conversation_handler) {
-    return;
-  }
-
-  DVLOG(1) << "Renamed conversation " << id << " to '" << new_name << "'";
-  OnConversationTitleChanged(conversation_handler, new_name);
+  ;
+  OnConversationTitleChanged(id, new_name);
 }
 
 void AIChatService::OnPremiumStatusReceived(GetPremiumStatusCallback callback,
@@ -809,12 +804,17 @@ void AIChatService::OnClientConnectionChanged(ConversationHandler* handler) {
   MaybeUnloadConversation(handler);
 }
 
-void AIChatService::OnConversationTitleChanged(ConversationHandler* handler,
-                                               std::string title) {
-  auto conversation_it = conversations_.find(handler->get_conversation_uuid());
-  CHECK(conversation_it != conversations_.end());
-  auto& conversation = conversation_it->second;
-  conversation->title = title;
+void AIChatService::OnConversationTitleChanged(
+    const std::string& conversation_uuid,
+    const std::string& new_title) {
+  auto conversation_it = conversations_.find(conversation_uuid);
+  if (conversation_it == conversations_.end()) {
+    DLOG(ERROR) << "Conversation not found for title change";
+    return;
+  }
+
+  auto& conversation_metadata = conversation_it->second;
+  conversation_metadata->title = new_title;
 
   OnConversationListChanged();
 
@@ -822,7 +822,7 @@ void AIChatService::OnConversationTitleChanged(ConversationHandler* handler,
   if (ai_chat_db_) {
     ai_chat_db_
         .AsyncCall(base::IgnoreResult(&AIChatDatabase::UpdateConversationTitle))
-        .WithArgs(handler->get_conversation_uuid(), std::move(title));
+        .WithArgs(conversation_uuid, new_title);
   }
 }
 
diff --git a/components/ai_chat/core/browser/ai_chat_service.h b/components/ai_chat/core/browser/ai_chat_service.h
index dfad1e2f8519..f8cf121389f7 100644
--- a/components/ai_chat/core/browser/ai_chat_service.h
+++ b/components/ai_chat/core/browser/ai_chat_service.h
@@ -93,8 +93,8 @@ class AIChatService : public KeyedService,
   void OnConversationEntryRemoved(ConversationHandler* handler,
                                   std::string entry_uuid) override;
   void OnClientConnectionChanged(ConversationHandler* handler) override;
-  void OnConversationTitleChanged(ConversationHandler* handler,
-                                  std::string title) override;
+  void OnConversationTitleChanged(const std::string& conversation_uuid,
+                                  const std::string& title) override;
   void OnAssociatedContentDestroyed(ConversationHandler* handler,
                                     int content_id) override;
 
@@ -174,7 +174,8 @@ class AIChatService : public KeyedService,
 
  private:
   // Key is uuid
-  using ConversationMap = std::map<std::string, mojom::ConversationPtr>;
+  using ConversationMap =
+      std::map<std::string, mojom::ConversationPtr, std::less<>>;
   using ConversationMapCallback = base::OnceCallback<void(ConversationMap&)>;
 
   void MaybeInitStorage();
diff --git a/components/ai_chat/core/browser/conversation_handler.cc b/components/ai_chat/core/browser/conversation_handler.cc
index 898b40671860..bbaf95bbf225 100644
--- a/components/ai_chat/core/browser/conversation_handler.cc
+++ b/components/ai_chat/core/browser/conversation_handler.cc
@@ -1685,7 +1685,7 @@ void ConversationHandler::OnClientConnectionChanged() {
 
 void ConversationHandler::OnConversationTitleChanged(std::string title) {
   for (auto& observer : observers_) {
-    observer.OnConversationTitleChanged(this, title);
+    observer.OnConversationTitleChanged(metadata_->uuid, title);
   }
 }
 
diff --git a/components/ai_chat/core/browser/conversation_handler.h b/components/ai_chat/core/browser/conversation_handler.h
index 8dfcd8a9d99b..9607c3ba6955 100644
--- a/components/ai_chat/core/browser/conversation_handler.h
+++ b/components/ai_chat/core/browser/conversation_handler.h
@@ -166,8 +166,9 @@ class ConversationHandler : public mojom::ConversationHandler,
 
     // Called when a mojo client connects or disconnects
     virtual void OnClientConnectionChanged(ConversationHandler* handler) {}
-    virtual void OnConversationTitleChanged(ConversationHandler* handler,
-                                            std::string title) {}
+    virtual void OnConversationTitleChanged(
+        const std::string& conversation_uuid,
+        const std::string& title) {}
     virtual void OnSelectedLanguageChanged(
         ConversationHandler* handler,
         const std::string& selected_language) {}