From ea278ee6eb3d4066ea1c1c42c95540e9a254efbf Mon Sep 17 00:00:00 2001
From: cogwheel0 <172976095+cogwheel0@users.noreply.github.com>
Date: Thu, 21 Aug 2025 16:19:21 +0530
Subject: [PATCH] fix: regenerate image gen

---
 lib/features/chat/providers/chat_providers.dart |  2 +-
 lib/features/chat/views/chat_page.dart          | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/lib/features/chat/providers/chat_providers.dart b/lib/features/chat/providers/chat_providers.dart
index ef05973..f601dd6 100644
--- a/lib/features/chat/providers/chat_providers.dart
+++ b/lib/features/chat/providers/chat_providers.dart
@@ -2313,7 +2313,7 @@ Future<void> cloneConversation(WidgetRef ref, String conversationId) async {
 }
 
 // Regenerate last message
-final regenerateLastMessageProvider = Provider<void Function()>((ref) {
+final regenerateLastMessageProvider = Provider<Future<void> Function()>((ref) {
   return () async {
     final messages = ref.read(chatMessagesProvider);
     if (messages.length < 2) return;
diff --git a/lib/features/chat/views/chat_page.dart b/lib/features/chat/views/chat_page.dart
index 1b1ae8f..15b5e34 100644
--- a/lib/features/chat/views/chat_page.dart
+++ b/lib/features/chat/views/chat_page.dart
@@ -934,6 +934,16 @@ class _ChatPageState extends ConsumerState<ChatPage> {
     }
 
     try {
+      // If assistant message has generated images and it's the last message,
+      // use image-only regenerate flow instead of text SSE regeneration
+      if (message.role == 'assistant' &&
+          (message.files?.any((f) => f['type'] == 'image') == true) &&
+          messageIndex == messages.length - 1) {
+        final regenerateImages = ref.read(regenerateLastMessageProvider);
+        await regenerateImages();
+        return;
+      }
+
       // Remove the assistant message we want to regenerate
       ref.read(chatMessagesProvider.notifier).removeLastMessage();